This repository hosts the official Jax implementation of the paper "NeRF on-the-go: Exploiting Uncertainty for Distractor-free NeRFs in the Wild" (CVPR 2024). For more details, please visit our project webpage.

This Repo is built upon Multinerf codebase.

# Clone the repo.
git clone https://github.com/cvg/nerf-on-the-go
cd nerf-on-the-go

# Make a conda environment.
conda create --name on-the-go python=3.9
conda activate on-the-go

# Prepare pip.
conda install pip
pip install --upgrade pip


# Install requirements.
pip install -r requirements.txt

# Manually install rmbrualla's `pycolmap` (don't use pip's! It's different).
git clone https://github.com/rmbrualla/pycolmap.git ./internal/pycolmap

# Confirm that all the unit tests pass.
./scripts/run_all_unit_tests.sh

You'll also need to update your JAX installation to support GPUs or TPUs.

pip install  jax==0.4.26 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install  jaxlib==0.4.26+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

srun -n 4 --mem-per-cpu=12000 --gpus=rtx_3090:1 --gres=gpumem:20g --time=4:00:00 --pty bash
conda activate on-the-go
module load eth_proxy gcc/8.2.0 cuda/12.1.1 cudnn/8.9.2.26

# Installs the wheel compatible with CUDA 12 and cuDNN 8.9 or newer.
pip install  jax==0.4.26 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install  jaxlib==0.4.26+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

The NVIDIA driver's CUDA version is 12.0 which is older than the ptxas CUDA version (12.4.131). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.

To download the "On-the-go" dataset, execute the following command:

bash ./scripts/download_on-the-go.sh

This script not only downloads the dataset but also downsamples the images as required.

For extracting features using the DINOv2, use the command below:

bash ./scripts/feature_extract.sh

After feature extraction, the dataset should be organized as

on-the-go
├── arcdetriomphe
│   ├── images
│   ├── images_{DOWNSAMPLE_RATE}
│   ├── features_{DOWNSAMPLE_RATE}
│   ├── split.json
│   ├── transforms.json
├── ....
│
└── tree
    ├── images_{DOWNSAMPLE_RATE}
    ├── ....
    └── transforms.json

• split.json: This file outlines the train and evaluation splits, following the naming conventions used in the RobustNeRF dataset, categorized as 'clutter' and 'clean'.
• transforms.json: Contains pose and intrinsic information, formatted according to the Blender dataset format, derived from COLMAP files. Refer to the Instant-NGP script for more details.

We plan to expand support to include custom datasets in future updates.

Example scripts for training, evaluating, and rendering can be found in scripts/. You'll need to change the paths to point to wherever the datasets are located. Gin configuration files for our model and some ablations can be found in configs/.

1. Training on-the-go:

bash scripts/train_on-the-go.sh 

2. Evaluating on-the-go:

bash scripts/eval_on-the-go.sh

3. Rendering on-the-go:

bash scirpts/render_on-the-go.sh

Tensorboard is supported for logging.

Since we use a different recording device for arc de triomphe and patio scene, the image downsample rate(4 instead of 8) and feature downsample rate(2 instead of 4) is different. Please use a separate script to train them by

bash scripts/train_on-the-go_HD.sh

About 80G gpu memory is needed to run current version.You may need to reduce the batch size (Config.batch_size) to avoid out of memory errors. If you do this, but want to preserve quality, be sure to increase the number of training iterations and decrease the learning rate by whatever scale factor you decrease batch size by.

• Custom dataset tutorial
• Support LPIPS calculation

If you use NeRF on-the-go, please cite

@InProceedings{Ren2024NeRF,
    title={NeRF on-the-go: Exploiting Uncertainty for Distractor-free NeRFs in the Wild},
    author={Ren, Weining and Zhu, Zihan and Sun, Boyang and Chen, Jiaqi and Pollefeys, Marc and Peng, Songyou},
    booktitle={IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2024}
}

Also, this code is built upon multinerf, feel free to cite this entire codebase as:

@misc{multinerf2022,
      title={{MultiNeRF}: {A} {Code} {Release} for {Mip-NeRF} 360, {Ref-NeRF}, and {RawNeRF}},
      author={Ben Mildenhall and Dor Verbin and Pratul P. Srinivasan and Peter Hedman and Ricardo Martin-Brualla and Jonathan T. Barron},
      year={2022},
      url={https://github.com/google-research/multinerf},
}

If there is any problem, please contact Weining by ren.weining@connect.hku.hk