PyTorch mip-NeRF

A reimplementation of mip-NeRF in PyTorch.

Not exactly 1-to-1 with the official repo, as we organized the code to out own liking (mostly how the datasets are structued, and hyperparam changes to run the code on a consumer level graphics card), made it more modular, and removed some repetitive code, but it achieves the same results.

Features

• Can use Spherical, or Spiral poses to generate videos for all 3 datasets
• Spherical:

video.mp4

• Spiral:

spiral.mp4

• Depth and Normals video renderings:
• Depth:

depth.mp4

• Normals:

normals.mp4

• Can extract meshes

mesh_.mp4

mesh.mp4

Future Plans

In the future we plan on implementing/changing:

• Factoring out more repetitive/redundant code, optimize gpu memory and rps
• Clean up and expand mesh extraction code
• Zoomed poses for multicam dataset
• Mip-NeRF 360: Unbounded Anti-Aliased Neural Radiance Fields support
• NeRV: Neural Reflectance and Visibility Fields for Relighting and View Synthesis support

Installation/Running

1. Create a conda environment using mipNeRF.yml
2. Get the training data
   1. run bash scripts/download_data.sh to download all 3 datasets: LLFF, Blender, and Multicam.
   2. Individually run the bash script corresponding to an individual dataset
      • bash scripts/download_llff.sh to download LLFF
      • bash scripts/download_blender.sh to download Blender
      • bash scripts/download_multicam.sh to download Multicam (Note this will also download the blender dataset since it's derived from it)
3. Optionally change config parameters: can change default parameters in config.py or specify with command line arguments
   • Default config setup to run on a high-end consumer level graphics card (~8-12GB)
4. Run python train.py to train
   • python -m tensorboard.main --logdir=log to start the tensorboard
5. Run python visualize.py to render a video from the trained model
6. Run python extract_mesh.py to render a mesh from the trained model

Code Structure

I explain the specifics of the code more in detail here but here is a basic rundown.

• config.py: Specifies hyperparameters.
• datasets.py: Base generic Dataset class + 3 default dataset implementations.
  • NeRFDataset: Base class that all datasets should inherent from.
  • Multicam: Used for multicam data as in the original mip-NeRF paper.
  • Blender: Used for the synthetic dataset as in original NeRF.
  • LLFF: Used for the llff dataset as in the original NeRF.
• loss.py: mip-NeRF loss, pretty much just MSE, but also calculates psnr.
• model.py: mip-NeRF model, not as modular as the way the original authors wrote it, but easier to understand its structure when laid out verbatim like this.
• pose_utils.py: Various functions used to generate poses.
• ray_utils.py: Various functions related involving rays that the model uses as input, most are used within the forward function of the model.
• scheduler.py: mip-NeRF learning rate scheduler.
• train.py: Trains a mip-NeRF model.
• visualize.py: Creates the videos using a trained mip-NeRF.

mip-NeRF Summary

Here's a summary on how NeRF and mip-NeRF work that I wrote when writing this originally.

• Summary

Results

_{^{All PSNRs are average PSNR (coarse + fine).}}

LLFF - Trex

Video:

video_.mp4

Depth:

depth_.mp4

Normals:

normals_.mp4

Blender - Lego

Video:

video.mp4

Depth:

depth.mp4

Normals:

normals.mp4

Multicam - Mic

Video:

video_.mp4

Depth:

depth_.mp4

Normals:

normals.mp4

References/Contributions

• Thanks to Nina for helping with the code
• Original NeRF Code in Tensorflow
• NeRF Project Page
• NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis
• Original mip-NeRF Code in JAX
• mip-NeRF Project Page
• Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields
• nerf_pl