PyTorch implementation of "3D-Aware Hypothesis & Verification for Generalizable Relative Object Pose Estimation" (ICLR 2024)

[project page]         [paper]

Setup Dependencies

conda create -n 3DAHV python=3.9
conda activate 3DAHV
bash ./install.sh

Experiments on CO3D

Data Preparation

Please refer to the instructions provided in RelPose++ for downloading and preprocessing CO3D. If necessary, you may need to adjust the values for ["CO3D"]["CO3D_DIR"] and ["CO3D"]["CO3D_ANNOTATION_DIR"] in the config.yaml file to match the actual directory path of your data.

Test pretrained model

We provide a model pretrained on the training set of CO3D. Please download it here. We store this pretrained model at ./models/Co3d_3DHAV by default. Run the following evaluation to get the results:

python ./test_co3d.py

Notably, the reproduced results might be slightly different from those reported in the paper. This is because the image pairs during testing are randomly sampled in the RelPose++ implementation.

Train model on CO3D

Run the following script to train the model:

python ./train_estimator_co3d.py

In our experiments, we trained the model on two A100 GPUs, and it took about two days to complete the training process.

Experiments on Objaverse

Data Preparation

Download the synthetic images generated by Zero123:

wget https://tri-ml-public.s3.amazonaws.com/datasets/views_release.tar.gz

Download images from COCO 2017 training set as background.

Modify ["OBJAVERSE"]["DATA_PATH"] and ["OBJAVERSE"]["COCO_IMAGE_ROOT"] in the config.yaml file according to the actual directory path of your data.

Train model on Objaverse

Run the following script to train the model:

python ./train_estimator_objaverse.py

Once the training is done, the evaluation is performed by running:

python ./test_objaverse.py

Experimentss on LINEMOD

Data Preparation

Download the LINEMOD dataset. We use the implementation in our previous work to preprocess the data.

Train model on LINEMOD

Run the following script to train the model:

python ./train_estimator_linemod.py

In our implentation, we finetune the model pretrained on Objaverse by default. We also tried training our model on Objaverse and LINEMOD at the same time, and we got the results comparable to those reported in our paper.

Once the training is done, the evaluation is performed by running:

python ./test_linemod.py

Citation

If you find the project useful, please consider citing:

@article{zhao20233d,
    title={3D-Aware Hypothesis \& Verification for Generalizable Relative Object Pose Estimation},
    author={Zhao, Chen and Zhang, Tong and Salzmann, Mathieu},
    journal={Proceedings of the International Conference on Learning Representations},
    year={2024}
  }
}