[arXiv]
[Project Page]
[Bibtex]
This repository contains code for "Cameras as Rays: Pose Estimation via Ray Diffusion" (ICLR 2024).
If you're installing on Windows, first install Visual Studio C++ Desktop Development. Then, ensure CUDA and conda are all on PATH.
Find a version of Pytorch compatible with your CUDA version from the Pytorch website. Note: if you pip install xformers, then use the latest Pytorch version.
conda create -n raydiffusion python=3.10
conda activate raydiffusion
conda install pytorch==2.2.0 torchvision==0.16.1 torchaudio==2.1.1 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install --pre -U xformers
pip install -r requirements.txt
pip install timm
Then, follow the directions to install Pytorch3D here.
git clone https://github.com/facebookresearch/pytorch3d.git
cd pytorch3d
python setup.py install # if this gives an issue, run in x64 Native Tools CP
Download the model weights from Google Drive.
Run ray diffusion with known bounding boxes (provided as a json):
python demo.py --model_dir models/co3d_diffusion --image_dir examples/robot/images \
--bbox_path examples/robot/bboxes.json --output_path robot.html
Run ray diffusion with bounding boxes extracted automatically from masks:
python demo.py --model_dir models/co3d_diffusion --image_dir examples/robot/images \
--mask_dir examples/robot/masks --output_path robot.html
Run ray regression:
python demo.py --model_dir models/co3d_regression --image_dir examples/robot/images \
--bbox_path examples/robot/bboxes.json --output_path robot.html
- Demo Code
- Evaluation Code
- Training Code
If you find this code helpful, please cite:
@InProceedings{zhang2024raydiffusion,
title={Cameras as Rays: Pose Estimation via Ray Diffusion},
author={Zhang, Jason Y and Lin, Amy and Kumar, Moneish and Yang, Tzu-Hsuan and Ramanan, Deva and Tulsiani, Shubham},
booktitle={International Conference on Learning Representations (ICLR)},
year={2024}
}