Skip to content

jonbarron/camp_zipnerf

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

19 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

CamP Zip-NeRF: A Code Release for CamP and Zip-NeRF

This is a branch of https://github.com/google-research/google-research/tree/master/camp_zipnerf, for those who don't want to have to check out that much larger repo.

This repository contains JAX code for two papers: Zip-NeRF and CamP. This is research code, and should be treated accordingly.

Setup

# Make a conda environment.
conda create --name camp_zipnerf python=3.11
conda activate camp_zipnerf

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

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

# Install camp_zipnerf.
git clone https://github.com/jonbarron/camp_zipnerf.git
cd camp_zipnerf
pip install -e .

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

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

Running

This code is designed to run on the NeRF Blender and Mip-NeRF 360 Datasets. If you want to onboard a new dataset, check out the instructions for the MultiNeRF code release, they should roughly apply here.

Running Zip-NeRF

To reproduce all of the Zip-NeRF results on the two datasets, just runscripts/zipnerf/all.sh --- be prepared to wait a long time! Unless you are at barron's workstation, you'll need to change some paths in the scripts that point to data folders. This script should recover a Zip-NeRF for all scenes, render out all test set images, render a few nice videos for the 360 dataset, and generate latex for tables. Because the paper's results were on 8 V100s and barron only has a GeForce RTX 3080 in his workstation, we dropped the batch size and learning rate by 8x while increasing the number of iterations by 8x (as per the linear scaling rule). Results are comparable to the paper (see the Zip-NeRF Results section below).

Note that we have results on the "single scale" mip-NeRF 360 Dataset (aka,"the mip-NeRF 360 dataset", see scripts/zipnerf/360_*.sh), as well as the "multi-scale" variant that Zip-NeRF presented in order to evaluate performance on multi-scale scenes, or scenes where the camera may be very far away (see scripts/zipnerf/ms360_*.sh). Unless you're explicitly interested in this multi-scale setting, you probably just want to use the "360_*.sh" scripts --- especially since they're the only ones that make nice videos.

Also note that the config we used in the Zip-NeRF paper to produce still images for evaluating error metrics (configs/zipnerf/360.gin) is different from the config we used to produce videos (configs/zipnerf/360_aglo128.gin). The latter includes the "affine GLO" per-image embedding that accounts for changes in exposure and white balance across images (see the supplement). Using this reduces floaters in the video renderings significantly, but also decreases test-set performance on these benchmarks. If you want maximally-accurate still images, use 360.gin, but if you want pretty videos, use 360_aglo128.gin.

Running CamP

There are multiple configs for CamP depending on the type of dataset (360 scene vs. object) and level of uncertainty (refining COLMAP poses vs. high noise).

  • camera_optim.gin contains the base CamP configuration.
  • camera_optim_perturbed.gin contains the CamP configuration for the perturbed 360 scene experiments. This will configure the data loader to add a significant amount of noise to the training camera parameters, and the training process will have to correct for this noise during scene fitting. Compared to be base config, the coarse-to-fine schedule is more aggressive to avoid falling into local minima.
  • camera_optim_perturbed_blender.gin contains the CamP configuration for the perturbed Blender scene experiments.
  • camera_optim_arkit.gin contains the CamP configuration for the ARKit experiments (note that these datasets have not been released at this time).

You can mix and match these configurations with the base NeRF configs (e.g., 360.gin). For example, to add camera refinement to Zip-NeRF, run the training pipeline like this:

python -m camp_zipnerf.train \
    --gin_configs=configs/zipnerf/360.gin \
    --gin_configs=configs/camp/camera_optim.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}/${SCENE}'" \
    --gin_bindings="Config.checkpoint_dir = '${CHECKPOINT_DIR}/${SCENE}'"

A script is available in scripts/camp/360_train.sh that runs Zip-NeRF + CamP on all of the Mip-NeRF 360 scenes.

To run different camera parameterizations, modify Config.camera_delta_cls to one of the following:

  • @camera_delta.FocalPoseCameraDelta: The Focal Pose parameterization.

  • @camera_delta.IntrinsicFocalPoseCameraDelta: The Focal Pose parameterization with radial distortion refinement.

  • @camera_delta.SE3CameraDelta: The SE3 parameterization, similar to BARF.

  • @camera_delta.SE3WithFocalCameraDelta: The SE3 parameterization with focal length refinement.

  • @camera_delta.IntrinsicSE3WithFocalCameraDelta: The SE3 parameterization with radial distortion and focal length refinement.

  • @camera_delta.SCNeRFCameraDelta: The SCNeRF parameterization.

    You can set these by modifying the config files, or by setting e.g., --gin_bindings="Config.config = '@camera_delta.FocalPoseCameraDelta'" in the launch command.

Zip-NeRF Results

For completeness we re-ran all results from the Zip-NeRF paper using this code repo, locally on barron's workstation. We've also linked to zip files of each rendered test-set, to allow others to evaluate their own error metrics.


Single-Scale mip-NeRF 360 Dataset Results:

Average:

PSNR SSIM
Code Release 28.56 0.828
Paper 28.54 0.828

Per-Scene

bicycle flowers garden stump treehill room counter kitchen bonsai
PSNR Code Release 25.85 22.33 28.22 27.35 23.95 33.04 29.12 32.36 34.79
Paper 25.80 22.40 28.20 27.55 23.89 32.65 29.38 32.50 34.46
SSIM Code Release 0.772 0.637 0.863 0.788 0.674 0.929 0.905 0.929 0.952
Paper 0.769 0.642 0.860 0.800 0.681 0.925 0.902 0.928 0.949

Download Results


Multi-Scale mip-NeRF 360 Dataset Results (Average):

PSNR 1x PSNR 2x PSNR 4x PSNR 8x SSIM 1x SSIM 2x SSIM 4x SSIM 8x
Code Release 28.15 29.90 31.43 32.38 0.822 0.890 0.931 0.952
Paper 28.25 30.00 31.57 32.52 0.822 0.892 0.933 0.954

Download Results


Single-Scale Per-Scene Blender Dataset Results:

chair drums ficus hotdog lego materials mic ship
PSNR Code Release 35.78 25.91 34.72 38.05 35.79 31.05 35.92 32.33
Paper 34.84 25.84 33.90 37.14 34.84 31.66 35.15 31.38
SSIM Code Release 0.987 0.948 0.987 0.987 0.983 0.968 0.992 0.937
Paper 0.983 0.944 0.985 0.984 0.980 0.969 0.991 0.929

Download Results


Citation

If you use this software package, please cite whichever constituent paper(s) you build upon, or feel free to cite this entire codebase as:

@misc{campzipnerf2024,
  title={{CamP Zip-NeRF}: {A} {Code} {Release} for {CamP} and {Zip-NeRF}},
  author={Jonathan T. Barron and Keunhong Park and Ben Mildenhall and John Flynn and Dor Verbin and Pratul Srinivasan and Peter Hedman and Philipp Henzler and Ricardo Martin-Brualla}
  year={2024},
  url={https://github.com/google-research/google-research/tree/master/camp_zipnerf},
}