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 演示。