Reloc3r 512

🚀 Reloc3r：用于可泛化、快速且准确视觉定位的相对相机位姿回归大规模训练

Reloc3r 是一个简单而有效的相机位姿估计框架，它将预训练的两视图相对相机位姿回归网络与多视图运动平均模块相结合。该框架在约 800 万个带位姿的图像对上进行训练，实现了惊人的性能和泛化能力，能够实时产生高质量的相机位姿估计。

相关链接

• 论文
• 代码

项目图示

🚀 快速开始

本项目提供了相机位姿估计的相关功能，下面将介绍如何安装、使用、评估和训练该项目。

📦 安装指南

1. 克隆 Reloc3r 仓库

git clone --recursive https://github.com/ffrivera0/reloc3r.git
cd reloc3r
# 如果你已经克隆了 reloc3r
# git submodule update --init --recursive

2. 使用 conda 创建环境

conda create -n reloc3r python=3.11 cmake=3.14.0
conda activate reloc3r 
conda install pytorch torchvision pytorch-cuda=12.1 -c pytorch -c nvidia  # 为你的系统使用正确的 cuda 版本
pip install -r requirements.txt
# 可选：你也可以安装额外的包来支持 HEIC 图像
pip install -r requirements_optional.txt

3. 可选：编译 RoPE 的 CUDA 内核

# Reloc3r 依赖于 RoPE 位置嵌入，你可以编译一些 CUDA 内核以加快运行时间
cd croco/models/curope/
python setup.py build_ext --inplace
cd ../../../

4. 可选：下载预训练模型

下载 Reloc3r-224 / Reloc3r-512 的检查点。运行下面的评估和演示代码时，预训练模型权重将自动下载。

💻 使用示例

相对位姿估计

使用 Reloc3r 可以估计你拍摄的图像和视频的相机位姿。对于相对位姿估计，尝试 wild_relpose.py 中的演示代码。我们提供了一些 图像对 用于演示。

# 用你的路径替换参数
python wild_relpose.py --v1_path ./data/wild_images/zurich0.jpg --v2_path ./data/wild_images/zurich1.jpg --output_folder ./data/wild_images/

可视化相对位姿

# 用你的路径替换参数
python visualization.py --mode relpose --pose_path ./data/wild_images/pose2to1.txt

视觉定位

对于视觉定位，wild_visloc.py 中的演示代码从自拍视频的采样帧中估计绝对相机位姿。

⚠️ 重要提示

该演示简单地使用第一帧和最后一帧作为数据库，这要求所有图像之间有重叠区域。此演示不支持线性运动。我们提供了一些 视频 作为示例。

# 用你的路径替换参数
python wild_visloc.py --video_path ./data/wild_video/ids.MOV --output_folder ./data/wild_video

可视化绝对位姿

# 用你的路径替换参数
python visualization.py --mode visloc --pose_folder ./data/wild_video/ids_poses/

📚 详细文档

相对相机位姿估计评估

要复现我们在 ScanNet1500 和 MegaDepth1500 上的评估，从 这里 下载数据集并解压到 ./data/。然后运行以下脚本，你将获得与我们论文中相似的结果。

bash scripts/eval_relpose.sh

💡 使用建议

要实现更快的推理速度，设置 --amp=1。这将启用 fp16 评估，在使用 Reloc3r-512 的 RTX 4090 上，速度将从 24 FPS 提高到 40 FPS，且不会损失任何精度。

视觉定位评估

要复现我们在 Cambridge 上的评估，从 这里 下载数据集并解压到 ./data/cambridge/。然后运行以下脚本，你将获得与我们论文中相似的结果。

bash scripts/eval_visloc.sh

训练

我们遵循 DUSt3R 来处理训练数据。下载数据集：CO3Dv2、ScanNet++、ARKitScenes、BlendedMVS、MegaDepth、DL3DV、RealEstate10K。

对于每个数据集，我们在 datasets_preprocess 目录中提供了一个预处理脚本，并在需要时提供了一个包含 图像对 列表的存档。你必须从官方来源自己下载数据集，同意其许可协议，并运行预处理脚本。

我们提供了一个在 RTX 3090 GPU 上使用 ScanNet++ 训练 Reloc3r 的示例脚本：

bash scripts/train_small.sh

要复现使用 8 个 H800 GPU 训练 Reloc3r-512 的过程，运行以下脚本：

bash scripts/train.sh

💡 使用建议

这些脚本与训练 Reloc3r 所使用的脚本并不严格相同，但应该足够接近。

📄 许可证

本项目未提及相关许可证信息。

📚 引用

如果你发现我们的工作对你的研究有帮助，请考虑引用：

@article{reloc3r,
  title={Reloc3r: Large-Scale Training of Relative Camera Pose Regression for Generalizable, Fast, and Accurate Visual Localization},
  author={Dong, Siyan and Wang, Shuzhe and Liu, Shaohui and Cai, Lulu and Fan, Qingnan and Kannala, Juho and Yang, Yanchao},
  journal={arXiv preprint arXiv:2412.08376},
  year={2024}
}

🙏 致谢

我们的实现基于以下优秀的仓库：

• Croco
• DUSt3R

感谢这些仓库的作者开源他们的代码。

📋 TODO 列表

• [x] 发布预训练权重和推理代码。
• [x] 发布 ScanNet1500、MegaDepth1500 和 Cambridge 数据集的评估代码。
• [x] 发布自拍图像和视频的示例代码。
• [x] 发布训练代码和数据。
• [ ] 发布其他数据集的评估代码。
• [ ] 发布视觉定位的加速版本。
• [ ] 发布 Gradio 演示。