HybridPose: 6D Object Pose Estimation under Hybrid Representations

This repository contains authors' implementation of HybridPose: 6D Object Pose Estimation under Hybrid Representations. Our implementation is based on PVNet. We warmly welcome any discussions related to our implementation and our paper. Please feel free to open an issue.

News (June 11, 2020): We have fixed a bug in Linemod dataset labeling as pointed out by JISock in this issue, and shanniruo in this issue. A small offset was incorrectly added to the translation vectors on Linemod dataset. Links to new labels and pre-trained weights are updated below. Our bug fix leads to marginal improvement of average ADD(-S) accuracy: from 0.9450447652 to 0.945402139. Thank you, JISock and shanniruo!

Introduction

HybridPose consists of intermediate representation prediction networks and a pose regression module. The prediction networks take an image as input, and output predicted keypoints, edge vectors, and symmetry correspondences. The pose regression module consists of a initialization sub-module and a refinement sub-module. The initialization sub-module solves a linear system with predicted intermediate representations to obtain an initial pose. The refinement sub-module utilizes GM robust norm to obtain the final pose prediction.

Download

git clone --recurse-submodules git@github.com:chensong1995/HybridPose.git

Environment set-up

Please install Anaconda first and execute the following commands:

conda create -y --name hybridpose python==3.7.4
conda install -y -q --name hybridpose -c pytorch -c anaconda -c conda-forge -c pypi --file requirements.txt
conda activate hybridpose

Compile the Ransac Voting Layer

The Ransac Voting Layer is used to generate keypoint coordinates from vector fields. Please execute the following commands (copied from PVNet):

cd lib/ransac_voting_gpu_layer
python setup.py build_ext --inplace

Compile the pose regressor

The pose regressor is written in C++ and has a Python wrapper. Please execute the following commands:

cd lib/regressor
make

Dataset set-up

We experimented HybridPose on Linemod and Occlusion Linemod. Let us first download the original datasets using the following commands:

python data/download_linemod.py
python data/download_occlusion.py

Let us then download our augumented labels to these two datasets. Our augumented labels include:

Keypoints: both 2D and 3D coordinates. These labels are generated using FSP.
Symmetry: Symmetry correspondences in 2D and the normal of symmetry plane in 3D. These labels are generated using SymSeg.
Segmentation masks: On Linemod, we create segmentation masks by projecting 3D models. On Occlusion Linemod, we use the segmentation masks provided in PVNet.

They are uploaded here:

Google Drive: Linemod, Occlusion Linemod.
Tencent Weiyun: Linemod, Occlusion Linemod.

The following commands unzip these labels to the correct directory:

unzip data/temp/linemod_labels.zip -d data/linemod
unzip data/temp/occlusion_labels.zip -d data/occlusion_linemod

Training

Please set the arguments in src/train_core.py execute the following command (note that we need to set LD_LIBRARY_PATH for the pose regressor):

# on bash shell
LD_LIBRARY_PATH=lib/regressor:$LD_LIBRARY_PATH python src/train_core.py
# on fish shell
env LD_LIBRARY_PATH="lib/regressor:$LD_LIBRARY_PATH" python src/train_core.py

If you use a different shell other than bash and fish, prepend "lib/regressor" to LD_LIBRARY_PATH and run python src/train_core.py.

Pre-trained weights

You can download our pre-trained weights below:

Linemod (Google Drive): ape, benchviseblue, cam, can, cat, driller, duck, eggbox, glue, holepuncher, iron, lamp, phone
Linemod (Tencent Weiyun): ape, benchviseblue, cam, can, cat, driller, duck, eggbox, glue, holepuncher, iron, lamp, phone
Occlusion Linemod (Google Drive): ape, can, cat, driller, duck, eggbox, glue, holepuncher
Occlusion Linemod (Tencent Weiyun): ape, can, cat, driller, duck, eggbox, glue, holepuncher

~~We have configured random seeds in src/train_core.py and expect you to re-produce identical weights by running our training script.~~ It turns out that completely reproducible results are not guaranteed across PyTorch releases, individual commits or different platforms. Furthermore, results need not be reproducible between CPU and GPU executions, even when using identical seeds. Our training uses two graphics cards with a batch size of 12.

After you download the pre-trained weights, unzip them somewhere and configure --load_dir in src/train_core.py to the unzipped weights (e.g. saved_weights/occlusion_linemod/ape/checkpoints/0.02/499).

Running src/train_core.py now will save both ground truth and predicted poses to a directory called output.

Evaluation

To evaluate ADD(-S) accuracy of predicted poses, please set the arguments in src/evaluate.py and run

python src/evaluate.py

Citation

If you find our work useful in your research, please kindly make a citation using:

@misc{song2020hybridpose,
    title={HybridPose: 6D Object Pose Estimation under Hybrid Representations},
    author={Chen Song and Jiaru Song and Qixing Huang},
    year={2020},
    eprint={2001.01869},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}

jakobamb/HybridPose