PyTorch implementation of "3D-Aware Hypothesis & Verification for Generalizable Relative Object Pose Estimation" (ICLR 2024)
[project page] [paper]
conda create -n 3DAHV python=3.9
conda activate 3DAHV
bash ./install.sh
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.
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.
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.
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.
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
Download the LINEMOD dataset. We use the implementation in our previous work to preprocess the data.
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
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}
}
}