MH6D: Multi-Hypothesis Consistency Learning for Category-Level 6D Object Pose Estimation

This is the PyTorch implementation of paper MH6D published in IEEE TNNLS by J. Liu, W. Sun, C. Liu, H. Yang, X. Zhang, and A. Mian.

Installation

Our code has been tested with

Ubuntu 20.04
Python 3.8.0
CUDA 11.0.3
PyTorch 1.8.0
gorilla-core 0.2.5.3

Some dependent packages：

gorilla

pip install gorilla-core==0.2.5.3

pointnet2

cd model/pointnet2
python setup.py install

Dataset

Download camera_train, camera_val, real_train, real_test, ground-truth annotations, mesh models, and segmentation results. Unzip and organize these files in ../data as follows:

data
├── CAMERA
│   ├── train
│   └── val
├── camera_full_depths
│   ├── train
│   └── val
├── Real
│   ├── train
│   └── test
├── gts
│   ├── val
│   └── real_test
├── obj_models
│   ├── train
│   ├── val
│   ├── real_train
│   └── real_test
├── segmentation_results
│   ├── train_trainedwoMask
│   ├── test_trainedwoMask
│   └── test_trainedwithMask
└── mean_shapes.npy

Run python scripts to prepare the datasets:

python data_processing.py

Evaluation

You can download our pretrained model epoch_30.pth and put it in the '../log1/3MH_sys' directory. Then, you can have a quick evaluation using the following command.

python test.py --config config1/3MH_sys.yaml

Note that there is a small mistake in the original evaluation code of NOCS for the 3D IOU metrics. We thank CATRE and SSC-6D for pointing out this. We have revised it and recalculated the metrics of some methods. The revised evaluation code is given in our released code.

Train

To train the model, remember to download the complete dataset, and organize and preprocess the dataset properly.

train_sys.py is the main file for training. You can start training using the following command.

python train_sys.py --gpus 0 --config config1/3MH_sys.yaml

Citation

If you find the code useful, please cite our paper.

@article{MH6D,
  author={Liu, Jian and Sun, Wei and Liu, Chongpei and Yang Hui and Zhang, Xing and Mian, Ajmal},
  journal={IEEE Transactions on Neural Networks and Learning Systems},
  title={MH6D: Multi-Hypothesis Consistency Learning for Category-Level 6D Object Pose Estimation},
  year={2024},
  publisher={IEEE},
  doi={10.1109/TNNLS.2024.3360712}
}

Acknowledgment

Our code is developed based on DPDN. We thank the authors for releasing the codes.

Licence

This project is licensed under the terms of the MIT license.

CNJianLiu/MH6D