jaxngp

This repository contains JAX implementations of: * a multiresolution hash encoder (JAX) * an accelerated volume renderer for fast training of NeRFs (CUDA + JAX), with * occupancy grid pruning during ray marching * early stop during ray color integration * an inference-time renderer for real-time rendering of NeRFs (CUDA + JAX) * a GUI for visualizing & interacting & exploring NeRFs @seimeicyx

Benchmarks

NeRF-synthetic

| | mic | ficus | chair | hotdog | materials | drums | ship | lego | average | |:--- |:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:| | @33.7k steps (this codebase) | 37.04 | 33.14 | 35.10 | 37.21 | 29.50 | 25.85 | 30.93 | 35.95 | 33.09 | | @51.2k steps (this codebase) | 37.07 | 33.17 | 35.16 | 37.26 | 29.50 | 25.86 | 30.94 | 36.03 | 33.124 | | paper (instant-ngp) | 36.22 | 33.51 | 35.00 | 37.40 | 29.78 | 26.02 | 31.10 | 36.39 | 33.176 |

^{For each scene, the network is trained on 100 training images (800x800 each) for 30k steps with default parameters, reported PSNR is averaged across 200 test images.}

Environment Setup

jaxngp manages environments with Nix, but it's also possible to setup the environment with any other package manager (e.g. Conda).

With Nix (recommended)

1. Install Nix with the official installer or the nix-installer.
2. With the nix executable available, clone this repository and setup environment: bash $ git clone https://github.com/blurgyy/jaxngp.git $ cd jaxngp/ $ NIXPKGS_ALLOW_UNFREE=1 nix develop --impure This will download (or build if necessary) all the dependencies, and opens a new shell with all the dependencies configured. > Note: to avoid the built environment being garbage collected when nix gc or nix-collect-garbage is called, append a --profile <PATH> argument: > bash > $ NIXPKGS_ALLOW_UNFREE=1 nix develop --impure --profile .git/devshell.profile >

With Conda

TODO

Running

The program's entrance is at python3 -m app.nerf. It provides three subcommands: train, test, and gui. Pass -h|--help to any of the subcommand to see its usage, e.g.:

Examples

• train
  • Just train with the default parameters on the lego scene from the NeRF-synthetic dataset: bash $ python3 -m app.nerf train data/nerf_synthetic/lego/transforms_train.json --exp-dir=logs/lego
  • Train for 10 epochs, with a batch size of 262144, on all the 400 (100*train + 100*validation + 200*test) images from the lego scene: bash $ python3 -m app.nerf train data/nerf_synthetic/lego --exp-dir=logs/lego-trainvaltest --{n-epochs=10,bs=262144}
  • Train on the training and validation splits of the drums scene, with a weight of 1e-5 on the Total Variation (TV) loss (by default this weight is 0): bash $ python3 -m app.nerf train data/nerf_synthetic/drums/transforms_{train,val}.json --exp-dir=logs/drums-trainval --tv-scale=1e-5
  • Train on the training split of the mic scene, validate with the validation split, validate after every epoch: bash $ python3 -m app.nerf train data/nerf_synthetic/mic/transforms_train.json --frames-val=data/nerf_synthetic/mic/transforms_val.json --exp-dir=logs/mic --validate-every=1 > Note: The validated images are logged to tensorboard, located under --exp-dir's logs/ directory. View it in browser with: > bash > $ tensorboard serve --logdir logs/mic/logs/ --bind_all > TensorBoard 2.10.0 at http://localhost:6006/ (Press CTRL+C to quit) >
• test
  • Test using the latest checkpoint under logs/mic directory, with the camera intrinsics and extrinsics of the mic scene's test split bash $ python3 -m app.nerf test data/nerf_synthetic/mic/transforms_test.json --ckpt=logs/mic/ --exp-dir=output
  • Test on the mic scene with given camera extrinsics, but override the camera's resolution to 1920x1080, and use white as background color: bash $ python3 -m app.nerf test data/nerf_synthetic/mic/transforms_test.json --ckpt=logs/mic/ --exp-dir=output --camera-override.{width=1920,height=1080} --render.bg 1 1 1
  • Test with a generated orbiting trajectory (see Demos for an example) on the mic scene, with resolution 1920x1080: bash $ python3 -m app.nerf test data/nerf_synthetic/mic --trajectory=orbit --ckpt=logs/mic/ --exp-dir=output --camera-override.{width=1920,height=1080}
• gui > Note: The gui subcommand accepts all the parameters of the train subcommand, and additionally a --viewport parameter (but the default values of --viewport are sane enough to leave as-is).
  • Train on all the 400 images from the lego scene, with 1280x720 as the default size of the rendering viewport: bash $ python3 -m app.nerf gui data/nerf_synthetic/lego --exp-dir=logs/gui-lego --viewport.{W=1280,H=720}

Running on Custom Data

A helper CLI (just a colmap wrapper via pycolmap) is provided for creating an Instant-NGP-compatible scene from a casually captured video or a directory of images:

bash $ python3 -m utils create-scene --help # see full usage usage: ... $ $ # capture or download a video $ mkdir -p data/_src $ curl https://github.com/blurgyy/jaxngp/assets/44701880/022a7b3c-344d-418f-aba0-0ccb9bfeb374 -Lo data/_src/gundam.mp4 $ $ # create a scene from the video, set scene bound to 16, with a background color model $ python3 -m utils create-scene data/_src/gundam.mp4 --root-dir=data/gundam --matcher=Sequential --fps=5 --bound=16 --bg [...]

After the scene has been created, the rest (training/validating/testing) are the same:

bash $ # train on all the registered images $ python3 -m utils train data/gundam --exp-dir=logs/gundam [...] $ $ # Render novel views, with a resolution of 1920x1080, save results as images and a video (video shown below in the Demo section) $ python3 -m app.nerf test data/gundam --{ckpt,exp-dir}=logs/gundam --trajectory=orbit --camera-override.{width=1920,height=1080,no-distortion} --orbit.high=1 --save-as="video and images"

Demos

• Orbiting camera trajectory:

gundam-nvs

^{Source video courtesy of @lzyronrico | Rendering time: ~18min for 288 frames | Resolution: 1920x1080 | Hardware: 1 GTX 1080Ti}

• Interactive rendering in the GUI:

gui-teaser

^{Source video courtesy of @filiptronicek | Hardware: 1 GTX 1080Ti}