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

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

• 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

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

• 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 Yang Fu if you have any further questions. This repository is for academic research use only.

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

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