Self-Supervised Geometric Correspondence for Category-Level 6D Object Pose Estimation in the Wild

This repository is the official implementation for the paper:

Self-Supervised Geometric Correspondence for Category-Level 6D Object Pose Estimation in the Wild
Kaifeng Zhang, Yang Fu, Shubhankar Borse, Hong Cai, Fatih Porikli, Xiaolong Wang
ICLR, 2023
Project Page / ArXiv

Progress

• Training code
• Evaluation code
• Pretrained models
• More datasets

Environment Setup

PyTorch with CUDA support are required. Our code is tested on python 3.8, torch 1.10.0, CUDA 11.1, and RTX 3090.

We recommend installing the required packages in the following order to avoid potential version conflicts:

pip install torch==1.10.0+cu111 torchvision==0.11.0+cu111 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install absl-py tensorboard opencv-python setuptools==59.5.0 trimesh kornia fvcore iopath matplotlib wandb scikit-learn scipy
pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu111_pyt1100/download.html

If the third step fails to find a valid version of pytorch3d, an alternative approach is to download the .tar.bz2 package from the Anaconda website (version: pytorch3d-0.6.1-py38_cu111_pyt1100.tar.bz2), and install the package with:

conda install /path/to/xxx.tar.bz2

Then, git clone our repository, and build the required SoftRas package located in PROJECT_ROOT/third-party/softras:

cd third-party/softras
python setup.py install

Data Preparation

To reproduce our experiments, please download the raw Wild6D dataset, unzip, and add dataset paths to training and testing configurations, for example, for the laptop training set, the correct path assignment is

--dataset_path /path/to/wild6d/laptop/ --test_dataset_path /path/to/wild6d/test_set/laptop/

Training

Our model uses a pretrained DINO for correspondence learning. Download the pretrained models from here and put them in the PROJECT_ROOT/pretrain/ directory. Also create the PROJECT_ROOT/log/ directory for logging. You can also assign custom paths via configuration.

Run training with train.py. We have provided an example training script for the laptop category, run by the following command:

bash scripts/train.sh

Testing

Run testing with predict.py. We have provided an example testing script, run by the following, run by the following command:

bash scripts/predict.sh

The testing script also offers visualization options. For example, use --vis_pred --vis_bbox for bounding box visualization.

Pretrained Models

We provide the pretrained models on Wild6D dataset containing all 5 categories.

To use, download the checkpoints in this link. Indicate the checkpoint path with the --model_path argument in scripts/predict.sh.

Citation

@article{zhang2022self,
  title={Self-Supervised Geometric Correspondence for Category-Level 6D Object Pose Estimation in the Wild},
  author={Zhang, Kaifeng and Fu, Yang and Borse, Shubhankar and Cai, Hong and Porikli, Fatih and Wang, Xiaolong},
  journal={arXiv preprint arXiv:2210.07199},
  year={2022}
}

Acknowledgement

The skeleton of our code is borrowed from ViSER. Thank you for the great work. Also thanks to the following external repos: SoftRas, Zero-Shot-Pose, U-CMR, PerceptualSimilarity, Objectron.