Awesome
MH6D: Multi-Hypothesis Consistency Learning for Category-Level 6D Object Pose Estimation
This is the PyTorch implementation of paper MH6D published in <b>IEEE TNNLS</b> by <a href="https://cnjliu.github.io/">J. Liu</a>, W. Sun, C. Liu, H. Yang, X. Zhang, and <a href="https://ajmalsaeed.net/">A. Mian</a>.
<p align="center"> <img src="image/Fig0.jpg" alt="intro" width="100%"/> </p>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:
pip install gorilla-core==0.2.5.3
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.