Update README.md

README.md

@@ -33,15 +33,28 @@ Trained on approximately 8 million posed image pairs, <strong>Reloc3r</strong> a
 </p>
 <be>
 
+
+## Table of Contents
+
+- [TODO List](#todo-list)
+- [Installation](#installation)
+- [Usage](#usage)
+- [Evaluation on Relative Camera Pose Estimation](#evaluation-on-relative-camera-pose-estimation)
+- [Evaluation on Visual Localization](#evaluation-on-visual-localization)
+- [Training](#training)
+- [Citation](#citation)
+- [Acknowledgments](#acknowledgments)
+
+
 ## TODO List
 
-- [x] Release pre-trained weights and inference code.
-- [x] Release evaluation code for ScanNet1500, MegaDepth1500 and Cambridge datasets.
-- [x] Release demo code for wild images and videos.
-- [ ] Release evaluation code for other datasets.
-- [ ] Release the accelerated version for visual localization.
-- [ ] Release Gradio Demo.
-- [ ] Release training code and data.
+- [x] Release pre-trained weights and inference code. 
+- [x] Release evaluation code for ScanNet1500, MegaDepth1500 and Cambridge datasets. 
+- [x] Release sample code for self-captured images and videos.
+- [x] Release training code and data.
+- [ ] Release evaluation code for other datasets. 
+- [ ] Release the accelerated version for visual localization. 
+- [ ] Release Gradio Demo.  
 
 
 ## Installation
@@ -57,7 +70,7 @@ cd reloc3r
 2. Create the environment using conda
 ```bash
 conda create -n reloc3r python=3.11 cmake=3.14.0
-conda activate reloc3r
+conda activate reloc3r 
 conda install pytorch torchvision pytorch-cuda=12.1 -c pytorch -c nvidia  # use the correct version of cuda for your system
 pip install -r requirements.txt
 # optional: you can also install additional packages to:
@@ -65,7 +78,7 @@ pip install -r requirements.txt
 pip install -r requirements_optional.txt
 ```
 
-3. Optional: Compile the cuda kernels for RoPE
+3. Optional: Compile the cuda kernels for RoPE 
 ```bash
 # Reloc3r relies on RoPE positional embeddings for which you can compile some cuda kernels for faster runtime.
 cd croco/models/curope/
@@ -73,33 +86,14 @@ python setup.py build_ext --inplace
 cd ../../../
 ```
 
-4. Optional: Download the checkpoints [Reloc3r-224](https://huggingface.co/siyan824/reloc3r-224)/[Reloc3r-512](https://huggingface.co/siyan824/reloc3r-512). The pre-trained model weights will automatically download when running the evaluation and demo code below.
-
+4. Optional: Download the checkpoints [Reloc3r-224](https://huggingface.co/siyan824/reloc3r-224)/[Reloc3r-512](https://huggingface.co/siyan824/reloc3r-512). The pre-trained model weights will automatically download when running the evaluation and demo code below. 
 
-## Relative Pose Estimation on ScanNet1500 and MegaDepth1500
 
-Download the datasets [here](https://drive.google.com/drive/folders/16g--OfRHb26bT6DvOlj3xhwsb1kV58fT?usp=sharing) and unzip it to `./data/`.
-Then run the following script. You will obtain results similar to those presented in our paper.
-```bash
-bash scripts/eval_relpose.sh
-```
-<strong>Note:</strong> To achieve faster inference speed, set `--amp=1`. This enables evaluation with `fp16`, which increases speed from <strong>24 FPS</strong> to <strong>40 FPS</strong> on an RTX 4090 with Reloc3r-512, without any accuracy loss.
+## Usage
 
+Using Reloc3r, you can estimate camera poses for images and videos you captured. 
 
-## Visual Localization on Cambridge
-
-Download the dataset [here](https://drive.google.com/file/d/1XcJIVRMma4_IClJdRq6rwBKX3ZPet5az/view?usp=sharing) and unzip it to `./data/cambridge/`.
-Then run the following script. You will obtain results similar to those presented in our paper.
-```bash
-bash scripts/eval_visloc.sh
-```
-
-
-## Demo for Wild Images
-
-In the demos below, you can run Reloc3r on your own data.
-
-For relative pose estimation, try the demo code in `wild_relpose.py`. We provide some [image pairs](https://drive.google.com/drive/folders/1TmoSKrtxR50SlFoXOwC4a9aGr18h00yy?usp=sharing) used in our paper.
+For relative pose estimation, try the demo code in `wild_relpose.py`. We provide some [image pairs](https://drive.google.com/drive/folders/1TmoSKrtxR50SlFoXOwC4a9aGr18h00yy?usp=sharing) used in our paper.  
 
 ```bash
 # replace the args with your paths
@@ -112,9 +106,10 @@ Visualize the relative pose
 python visualization.py --mode relpose --pose_path ./data/wild_images/pose2to1.txt
 ```
 
-For visual localization, the demo code in `wild_visloc.py` estimates absolute camera poses from sampled frames in self-captured videos.
+For visual localization, the demo code in `wild_visloc.py` estimates absolute camera poses from sampled frames in self-captured videos. 
 
-<strong>Important</strong>: The demo uses the first and last frames as the database, which <strong>requires</strong> overlapping regions among all images. This demo does <strong>not</strong> support linear motion. We provide some [videos](https://drive.google.com/drive/folders/1sbXiXScts5OjESAfSZQwLrAQ5Dta1ibS?usp=sharing) as examples.
+> [!IMPORTANT]
+> The demo simply uses the first and last frames as the database, which <strong>requires</strong> overlapping regions among all images. This demo does <strong>not</strong> support linear motion. We provide some [videos](https://drive.google.com/drive/folders/1sbXiXScts5OjESAfSZQwLrAQ5Dta1ibS?usp=sharing) as examples. 
 
 ```bash
 # replace the args with your paths
@@ -128,9 +123,50 @@ python visualization.py --mode visloc --pose_folder ./data/wild_video/ids_poses/
 ```
 
 
+## Evaluation on Relative Camera Pose Estimation 
+
+To reproduce our evaluation on ScanNet1500 and MegaDepth1500, download the datasets [here](https://drive.google.com/drive/folders/16g--OfRHb26bT6DvOlj3xhwsb1kV58fT?usp=sharing) and unzip it to `./data/`.
+Then run the following script. You will obtain results similar to those presented in our paper.
+```bash
+bash scripts/eval_relpose.sh
+```
+
+> [!NOTE]
+> To achieve faster inference speed, set `--amp=1`. This enables evaluation with `fp16`, which increases speed from <strong>24 FPS</strong> to <strong>40 FPS</strong> on an RTX 4090 with Reloc3r-512, without any accuracy loss.
+
+
+## Evaluation on Visual Localization 
+
+To reproduce our evaluation on Cambridge, download the dataset [here](https://drive.google.com/file/d/1XcJIVRMma4_IClJdRq6rwBKX3ZPet5az/view?usp=sharing) and unzip it to `./data/cambridge/`.
+Then run the following script. You will obtain results similar to those presented in our paper.
+```bash
+bash scripts/eval_visloc.sh
+```
+
+
+## Training
+
+We follow [DUSt3R](https://github.com/naver/dust3r) to process the training data. Download the datasets: [CO3Dv2](https://github.com/facebookresearch/co3d), [ScanNet++](https://kaldir.vc.in.tum.de/scannetpp/), [ARKitScenes](https://github.com/apple/ARKitScenes), [BlendedMVS](https://github.com/YoYo000/BlendedMVS), [MegaDepth](https://www.cs.cornell.edu/projects/megadepth/), [DL3DV](https://dl3dv-10k.github.io/DL3DV-10K/), [RealEstate10K](https://google.github.io/realestate10k/). 
+
+For each dataset, we provide a preprocessing script in the `datasets_preprocess` directory and an archive containing the list of [pairs](https://drive.google.com/drive/folders/193Lv5YB-2OVkqK3k6vnZJi36-FRPcAuu?usp=sharing) when needed. You have to download the datasets yourself from their official sources, agree to their license, and run the preprocessing script.
+
+We provide a sample script to train Reloc3r with ScanNet++ on an RTX 3090 GPU
+```bash
+bash scripts/train_small.sh
+```
+
+To reproduce our training for Reloc3r-512 with 8 H800 GPUs, run the following script
+```bash
+bash scripts/train.sh
+```
+
+> [!NOTE]
+> They are not strictly equivalent to what was used to train Reloc3r, but they should be close enough.
+
+
 ## Citation
 
-If you find our work helpful in your research, please consider citing:
+If you find our work helpful in your research, please consider citing: 
 ```
 @article{reloc3r,
   title={Reloc3r: Large-Scale Training of Relative Camera Pose Regression for Generalizable, Fast, and Accurate Visual Localization},
@@ -148,4 +184,5 @@ Our implementation is based on several awesome repositories:
 - [Croco](https://github.com/naver/croco)
 - [DUSt3R](https://github.com/naver/dust3r)
 
-We thank the respective authors for open-sourcing their code.
+We thank the respective authors for open-sourcing their code.
+