Repository files navigation

This is the PyTorch implementation of paperMH6D published inIEEE TNNLS byJ. Liu, W. Sun, C. Liu, H. Yang, X. Zhang, andA. Mian.

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/pointnet2python setup.py install

Downloadcamera_train,camera_val,real_train,real_test,ground-truth annotations,mesh models, andsegmentation 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

You can download our pretrained modelepoch_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 ofNOCS for the 3D IOU metrics. We thankCATRE andSSC-6D for pointing out this. We have revised it and recalculated the metrics of some methods. The revised evaluation code is given in our releasedcode.

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

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}}

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

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