PyTorch mip-NeRF

A reimplementation of mip-NeRF in PyTorch.

nerfTomipnerf

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:

https://user-images.githubusercontent.com/42706447/171090423-2cf37b0d-44c9-4394-8c4a-46f19b0eb304.mp4

Spiral:

https://user-images.githubusercontent.com/42706447/171099856-f7340263-1d65-4fbe-81e7-3b2dfa9e93b8.mp4

Depth and Normals video renderings:
Depth:

https://user-images.githubusercontent.com/42706447/171091394-ce73822c-689f-496b-8821-78883e8b90d4.mp4

Normals:

https://user-images.githubusercontent.com/42706447/171091457-c795855e-f8f8-4515-ae62-7eeb707d17bc.mp4

Can extract meshes

https://user-images.githubusercontent.com/42706447/171100048-8f57fc9a-4be5-44c2-93dd-ee5f6b54dd6e.mp4

https://user-images.githubusercontent.com/42706447/171092108-b60130b5-297d-4e72-8d3a-3e5a29c83036.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

Create a conda environment using mipNeRF.yml
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)
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)
Run python train.py to train
- python -m tensorboard.main --logdir=log to start the tensorboard
Run python visualize.py to render a video from the trained model
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.