Category-Level 6D Object Pose Estimation in the Wild: A Semi-Supervised Learning Approach and A New Dataset

This repository contains the official implementation for the following paper:

Category-Level 6D Object Pose Estimation in the Wild: A Semi-Supervised Learning Approach and A New Dataset
Yang Fu, Xiaolong Wang
Project Page | Paper (arXiv)

NeurIPS 2022

Progress

• Pretrained models
• Evaluation script
• Model architecture
• Training script
• Code cleaning

Requirements

Environments

• Install PyTorch>=1.8.1 here based the package management tool you used and your cuda version (older PyTorch versions may work but have not been tested)
• pip install -r requirements.txt

Data

Download and unzip Wild6D data from Google Drive or OneDrive (Testing set is only needed for evaluation). We highly recommend you to download Wild6D data via gdown. For example, you can download the testing data with the following command.

gdown --id 1AWLX6V0kAyiTFdkGg4FrkEaQ5jbUVtaj

We also provide a script that allows downloading all dataset files at once. In order to do so, execute the download script,

bash tools/download.sh

Unzip and organize these files in $ROOT/data as the following structure:

data
├── Wild6D
│   ├── bottle
│   ├── bowl
│   ├── camera
│   ├── laptop
│   ├── mug
│   └── test_set
│       ├──pkl_annotations
│       │   ├── bottle
│       │   ├── bowl
│       │       ...
│       ├── bottle
│       ├── bowl
│           ...
├── meshes

Evaluation

• Download the pretrained model for different categories from Google Drive to $MODEL_PATH

• Pick a category name $CATEGORY_NAME which you can find definition in evaluate_wild6d.py. Using the following code to evaluate the pretrained model.

  python evaluate_wild6d.py --use_nocs_map --implicit --model $MODEL_PATH --select_class $CATEGORY_NAME
  

• Note that the results may be slightly different from the number reported in paper, since we further clean the dataset recently. We also provide the estimation results that align with the paper number as the reference.

Contact

Contact Yang Fu if you have any further questions. This repository is for academic research use only.

Acknowledgments

Our codebase builds heavily on NOCS and Shape-Prior. Thanks for open-sourcing!.

Citation

@inproceedings{fucategory,
  title={Category-Level 6D Object Pose Estimation in the Wild: A Semi-Supervised Learning Approach and A New Dataset},
  author={Fu, Yang and Wang, Xiaolong},
  booktitle={Advances in Neural Information Processing Systems},
  year={2022}
}