/ Machine Learning

6D Object Pose Estimation under Hybrid Representations

6D Object Pose Estimation under Hybrid Representations

HybridPose

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. Approach overview

Download

git clone --recurse-submodules [email protected]: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:

  • Blender meshes on Linemod objects: For some reasons, pose labels on Linemod are not aligned perfectly with the 3D models. After discussions with the authors of PVNet, we followed their advice and used Blender meshes to correct Linemod pose labels.
  • 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: 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:

We have configured random seeds in src/train_core.py and expect you to re-produce identical weights by running our training script. 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

GitHub