This is a re-implementation of the "Robust 6D Object Pose Estimation by Learning RGB-D Features" paper. The original implementation can be found on https://github.com/mentian/object-posenet.

The PSPNet under model/psp is from https://github.com/Lextal/pspnet-pytorch.
The transformation functions under utils/transformations.py (used for converting quaternions to matrices) are from https://github.com/cgohlke/transformations.

The pose estimation network is trained on the LINEMOD dataset. A preprocessed version of this dataset, that also contains object segmentations obtained by a SegNet, is available here. This preprocessed dataset is used during evaluation.

Our trained model can be downloaded here.

Install the requirements from requirements.txt, download the dataset and the trained model and adjust the paths to the dataset directory and to the trained model if necessary. The python version we used is Python 3.6.5.

Cuda is required for training. Adjust the dataset_path in train_linemod in train.py and then run train.py

Cuda is required for evaluation. Adjust the dataset_path and path_to_trained_model in evaluate.py and then run evaluate.py

We have encountered an issue with calling model.eval() after loading the trained model. When doing this, the accuracy of the predictions is nearly 0. It seems that this is a known issue (https://discuss.pytorch.org/t/performance-highly-degraded-when-eval-is-activated-in-the-test-phase/3323) investigate this further, which is why in the demo notebook and in evaluate.py we did not call model.eval().