This is the official PyTorch implementation of End-to-End Probabilistic Perspective-n-Points for 6DoF object pose estimation. [paper]

The code is modified from the official implementation of CDPN, and is used for benchmarking only. We will not maintain this code except for bug fixes.

EPro-PnP-6DoF reuses the off-the-shelf 6DoF pose estimation network CDPN. The original CDPN adopts two decoupled branches: a direct prediction branch for position, and a dense correspondence branch (PnP-based) for orientation. EPro-PnP-6DoF keeps only the dense correspondence branch (with minor modifications to the output layer for the 2-channel weight map), to which the EPro-PnP layer is appended for end-to-end 6DoF pose learning.

The code has been tested in the environment described as follows:

• Linux (tested on Ubuntu 16.04/18.04)
• Python 3.6
• PyTorch 1.5.0

An example script for installing the python dependencies under CUDA 10.2:

# Create conda environment
conda create -y -n epropnp_6dof python=3.6
conda activate epropnp_6dof
conda install -y pip

# Install pytorch
conda install pytorch==1.5.0 torchvision==0.6.0 cudatoolkit=10.2 -c pytorch

# Install other dependencies
pip install opencv-python==4.5.1.48 pyro-ppl==1.4.0 PyYAML==5.4.1 matplotlib termcolor plyfile easydict scipy progress numba tensorboardx

Please refer to this link for instructions. In case you have trouble downloading the LineMOD dataset, we have uploaded a copy here. Afterwards, the dataset folders should be structured as follows:

EPro-PnP-6DoF/
├── dataset/
│   ├── bg_images/
│   │   └── VOC2012/
│   └── lm/
│       ├── models/
│       │   ├── ape/
│       │   …
│       ├── imgn/
│       │   ├── ape/
│       │   …
│       ├── real_test/
│       │   ├── ape/
│       │   …
│       └── real_train/
│           ├── ape/
│           …
├── lib/    
├── tools/
…

Since the experiments in the main paper (models v1), we have heavily refactored the code of the EPro-PnP layer, resulting in slightly different numerical behavior. We release the models trained with the current refactored code (models v1b) at [Google Drive | Baidu Pan]. The results of both versions are shown below.

If you use the epropnp_cdpn_init* config, please download the checkpoint cdpn_stage_1.pth from [Google Drive | Baidu Pan], and move it to EPro-PnP-6DoF/checkpoints/cdpn_stage_1.pth.

To start training, enter the directory EPro-PnP-6DoF/tools, and run:

python main.py --cfg /PATH/TO/CONFIG  # configs are located in EPro-PnP-6DoF/tools/exp_cfg

By default GPU 0 is used, you can set the environment variable CUDA_VISIBLE_DEVICES to change this behavior.

Checkpoints, logs and visualizations will be saved to EPro-PnP-6DoF/exp. You can run TensorBoard to plot the logs:

tensorboard --logdir ../exp

To test and evaluate on the LineMOD test split, please edit the config file and

1. set the load_model option to the path of the checkpoint file,
2. change the test option from False to True.

After saving the test config, enter the directory EPro-PnP-6DoF/tools, and run:

python main.py --cfg /PATH/TO/CONFIG

Logs and visualizations will be saved to EPro-PnP-6DoF/exp.