将相机视为光线

[arXiv] [项目主页] [引用] [Colab]

本仓库包含"将相机视为光线：基于光线扩散的姿态估计"（ICLR 2024）的代码。

克隆仓库：

git clone --depth=1 --branch=main https://github.com/jasonyzhang/RayDiffusion.git

环境配置

我们建议使用conda环境来管理依赖项。从Pytorch官网安装与您的CUDA版本兼容的Pytorch版本。

conda create -n raydiffusion python=3.10
conda activate raydiffusion
conda install pytorch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 pytorch-cuda=11.8 -c pytorch -c nvidia
conda install xformers -c xformers
pip install -r requirements.txt

然后，按照这里的指示安装Pytorch3D。我们建议使用对应Python/Pytorch/CUDA版本的预构建轮子安装Pytorch3D：

pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu118_pyt211/download.html

如果使用预构建轮子安装遇到问题，您也可以尝试从源代码构建，但这将花费更长时间。

运行演示

从Google Drive下载模型权重。

gdown https://drive.google.com/uc\?id\=1anIKsm66zmDiFuo8Nmm1HupcitM6NY7e
unzip models.zip

使用已知边界框（以json格式提供）运行光线扩散：

python demo.py  --model_dir models/co3d_diffusion --image_dir examples/robot/images \
    --bbox_path examples/robot/bboxes.json --output_path robot.html

使用从掩码自动提取的边界框运行光线扩散：

python demo.py  --model_dir models/co3d_diffusion --image_dir examples/robot/images \
    --mask_dir examples/robot/masks --output_path robot.html

运行光线回归：

python demo.py  --model_dir models/co3d_regression --image_dir examples/robot/images \
    --bbox_path examples/robot/bboxes.json --output_path robot.html

训练

光线扩散的训练命令：

accelerate launch --multi_gpu --gpu_ids 0,1,2,3,4,5,6,7 --num_processes 8 train.py \
    training.batch_size=8 training.max_iterations=450000

有关训练的更详细说明，请参见docs/train.md。

评估

有关如何运行评估代码的说明，请参见docs/eval.md。

引用将相机视为光线

如果您发现这段代码有帮助，请引用：

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