PlenOctrees: NeRF-SH Training & Conversion

Overview

PlenOctrees Official Repo: NeRF-SH training and conversion

This repository contains code to train NeRF-SH and to extract the PlenOctree, constituting part of the code release for:

PlenOctrees for Real Time Rendering of Neural Radiance Fields
Alex Yu, Ruilong Li, Matthew Tancik, Hao Li, Ren Ng, Angjoo Kanazawa

https://alexyu.net/plenoctrees

Please see the following repository for our C++ PlenOctrees volume renderer: https://github.com/sxyu/volrend

Setup

Please use conda for a replicable environment.

conda env create -f environment.yml
conda activate plenoctree
pip install --upgrade pip

Or you can install the dependencies manually by:

conda install pytorch torchvision cudatoolkit=11.0 -c pytorch
conda install tqdm
pip install -r requirements.txt

[Optional] Install GPU and TPU support for Jax. This is useful for NeRF-SH training. Remember to change cuda110 to your CUDA version, e.g. cuda102 for CUDA 10.2.

pip install --upgrade jax jaxlib==0.1.65+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

NeRF-SH Training

We release our trained NeRF-SH models as well as converted plenoctrees at Google Drive. You can also use the following commands to reproduce the NeRF-SH models.

Training and evaluation on the NeRF-Synthetic dataset (Google Drive):

export DATA_ROOT=./data/NeRF/nerf_synthetic/
export CKPT_ROOT=./data/Plenoctree/checkpoints/syn_sh16/
export SCENE=chair
export CONFIG_FILE=nerf_sh/config/blender

python -m nerf_sh.train \
    --train_dir $CKPT_ROOT/$SCENE/ \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/

python -m nerf_sh.eval \
    --chunk 4096 \
    --train_dir $CKPT_ROOT/$SCENE/ \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/

Note for SCENE=mic, we adopt a warmup learning rate schedule (--lr_delay_steps 50000 --lr_delay_mult 0.01) to avoid unstable initialization.

Training and evaluation on TanksAndTemple dataset (Download Link) from the NSVF paper:

export DATA_ROOT=./data/TanksAndTemple/
export CKPT_ROOT=./data/Plenoctree/checkpoints/tt_sh25/
export SCENE=Barn
export CONFIG_FILE=nerf_sh/config/tt

python -m nerf_sh.train \
    --train_dir $CKPT_ROOT/$SCENE/ \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/

python -m nerf_sh.eval \
    --chunk 4096 \
    --train_dir $CKPT_ROOT/$SCENE/ \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/

PlenOctrees Conversion and Optimization

Before converting the NeRF-SH models into plenoctrees, you should already have the NeRF-SH models trained/downloaded and placed at ./data/PlenOctree/checkpoints/{syn_sh16, tt_sh25}/. Also make sure you have the training data placed at ./data/{NeRF/nerf_synthetic, TanksAndTemple}.

To reproduce our results in the paper, you can simplly run:

# NeRF-Synthetic dataset
python -m octree.task_manager octree/config/syn_sh16.json --gpus="0 1 2 3"

# TanksAndTemple dataset
python -m octree.task_manager octree/config/tt_sh25.json --gpus="0 1 2 3"

The above command will parallel all scenes in the dataset across the gpus you set. The json files contain dedicated hyper-parameters towards better performance (PSNR, SSIM, LPIPS). So in this setting, a 24GB GPU is needed for each scene and in averange the process takes about 15 minutes to finish. The converted plenoctree will be saved to ./data/PlenOctree/checkpoints/{syn_sh16, tt_sh25}/$SCENE/octrees/.

Below is a more straight-forward script for demonstration purpose:

export DATA_ROOT=./data/NeRF/nerf_synthetic/
export CKPT_ROOT=./data/PlenOctree/checkpoints/syn_sh16
export SCENE=chair
export CONFIG_FILE=nerf_sh/config/blender

python -m octree.extraction \
    --train_dir $CKPT_ROOT/$SCENE/ --is_jaxnerf_ckpt \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/ \
    --output $CKPT_ROOT/$SCENE/octrees/tree.npz

python -m octree.optimization \
    --input $CKPT_ROOT/$SCENE/tree.npz \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/ \
    --output $CKPT_ROOT/$SCENE/octrees/tree_opt.npz

python -m octree.evaluation \
    --input $CKPT_ROOT/$SCENE/octrees/tree_opt.npz \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/

# [Optional] Only used for in-browser viewing.
python -m octree.compression \
    $CKPT_ROOT/$SCENE/octrees/tree_opt.npz \
    --out_dir $CKPT_ROOT/$SCENE/ \
    --overwrite

MISC

Project Vanilla NeRF to PlenOctree

A vanilla trained NeRF can also be converted to a plenoctree for fast inference. To mimic the view-independency propertity as in a NeRF-SH model, we project the vanilla NeRF model to SH basis functions by sampling view directions for every points in the space. Though this makes converting vanilla NeRF to a plenoctree possible, the projection process inevitability loses the quality of the model, even with a large amount of sampling view directions (which takes hours to finish). So we recommend to just directly train a NeRF-SH model end-to-end.

Below is a example of projecting a trained vanilla NeRF model from JaxNeRF repo (Download Link) to a plenoctree. After extraction, you can optimize & evaluate & compress the plenoctree just like usual:

export DATA_ROOT=./data/NeRF/nerf_synthetic/ 
export CKPT_ROOT=./data/JaxNeRF/jaxnerf_models/blender/ 
export SCENE=drums
export CONFIG_FILE=nerf_sh/config/misc/proj

python -m octree.extraction \
    --train_dir $CKPT_ROOT/$SCENE/ --is_jaxnerf_ckpt \
    --config $CONFIG_FILE \
    --data_dir $DATA_ROOT/$SCENE/ \
    --output $CKPT_ROOT/$SCENE/octrees/tree.npz \
    --projection_samples 100 \
    --radius 1.3

Note --projection_samples controls how many sampling view directions are used. More sampling view directions give better projection quality but takes longer time to finish. For example, for the drums scene in the NeRF-Synthetic dataset, 100 / 10000 sampling view directions takes about 2 mins / 2 hours to finish the plenoctree extraction. It produce raw plenoctrees with PSNR=22.49 / 23.84 (before optimization). Note that extraction from a NeRF-SH model produce a raw plenoctree with PSNR=25.01.

Comments
  • error when run octree extraction

    error when run octree extraction

    Hello, thank you for providing codes! I met the problem when trying to extract the octree from a Jax checkpoint.

    I ran the command as: python -m octree.extraction --train_dir checkpoints/chair/ --is_jaxnerf_ckpt --config nerf_sh/config/blender --data_dir ../../data/3D/synthetic/chair --output checkpoints/chair/octrees/tree.npz

    and I got the error as: File "/workspace/NeRF/plenoctree/octree/extraction.py", line 184, in calculate_grid_weights opts = _C.RenderOptions() AttributeError: 'NoneType' object has no attribute 'RenderOptions'

    Could you please help me?

    opened by Holmes-Alan 4
  • Training problem

    Training problem

    When I Training and evaluation on TanksAndTemple dataset. python -m nerf_sh.train --train_dir $CKPT_ROOT/$SCENE/ --config$CONFIG_FILE --data_dir $DATA_ROOT/$SCENE/ (NO Install GPU and TPU support for Jax.) There are the following errors bddb316a42cc6668d2bb1fed23a97f6 7633f460bbda3cd6f3040d11fae3c34

    opened by SSground 4
  • Fatal Python error: Aborted when use library svox

    Fatal Python error: Aborted when use library svox

    Hi, when i run python -m octree.task_manager octree/config/tt_sh25.json --gpus="0" image This error rise in line 205 image But this library PyTorch CUDA Extension of PlenOctrees and installed by git url. How can I solve this problem?

    opened by katealenic 4
  • octree optimization does not work for LLFF dataset

    octree optimization does not work for LLFF dataset

    Hello, I found that the codes seems not to work for LLFF dataset. I can successfully train the model for scene synthesis. However, when I try to use octree to extract the nerf model, it cannot work. i have tested that the codes work well for blender dataset. I wonder if the authors have any solutions. Thanks in advance.

    opened by Holmes-Alan 3
  • error when run the training code

    error when run the training code

    I tried to install the code following the same procedure. However, I got the error when installing Jax, Jaxlib and Flax as follows.

    raise RuntimeError(msg) RuntimeError: jaxlib is version 0.1.57, but this version of jax requires version >= 0.3.7.

    I checked that I installed Jax==0.3.9 and Jaxlib==0.1.57, which meets the requirement. I do not know what the problem is. I am looking forward to your reply.

    opened by Holmes-Alan 1
  • RuntimeError when running Octree extraction

    RuntimeError when running Octree extraction

    Hi, i've tried this now for a while and keep encountering this issue.

    Training goes well, but when i run the extraction i get this:

    image

    Anyone know what is wrong? or how to resolve it?

    opened by KevCes 0
  • Error in svox.volume_render_image_backward : too many resources requested for launch

    Error in svox.volume_render_image_backward : too many resources requested for launch

    When I do the PlenOctrees Optimization with python -m octree.optimization, one error comes up:

    Error in svox.volume_render_image_backward : too many resources requested for launch
    

    which results in no change of the validation_psnr and then Stop since overfitting. May I ask what the problem is and how to solve it? Many thanks.

    opened by weihaosky 0
  • Invoked in extraction

    Invoked in extraction

    I run the script in MISC to Project Vanilla NeRF to PlenOctree an error happened

    image

    and I finally find out only running this line

    https://github.com/sxyu/plenoctree/blob/f0b82631199a1aa7dc9ce263c08980a3a7504014/octree/nerf/models.py#L70

    import this package will get that error.

    Anyone knows what happened?

    opened by zhywanna 1
  • How to extract SHs from vanilla TF-NeRF or pyTorch NeRF

    How to extract SHs from vanilla TF-NeRF or pyTorch NeRF

    @sxyu Could the authors please answer this question about "Which implementation can be used to perform extraction without is_jaxnerf_ckpt flag" ?

    https://github.com/sxyu/plenoctree/issues/6#issue-925605838

    opened by zhywanna 0
  • how do you visualize the npz files by nerfvis?

    how do you visualize the npz files by nerfvis?

    I understand that you release the trained NeRF-SH models as well as converted plenoctrees. However, when I used your converted plenoctrees for 3D visualization like using "ship.npz" from your file as following,

    import nerfvis scene = nerfvis.Scene("My 3D model") scene.export(dirname='checkpoints/ship', display=True, tree_file='tree.npz', url="100.80.243.129", port=80)

    I cannot get a 3D view from my browser. It always show loading at 100% but it never visualize the view, like image

    I assume that the problem is from the plenoctree compression step but I still cannot get it resolved. Could you please help me out? Thank you!

    opened by Holmes-Alan 0
Owner
Alex Yu
Undergrad at UC Berkeley
Alex Yu
Matlab Python Heuristic Battery Opt - SMOP conversion and manual conversion

SMOP is Small Matlab and Octave to Python compiler. SMOP translates matlab to py

Tom Xu 1 Jan 12, 2022
(Arxiv 2021) NeRF--: Neural Radiance Fields Without Known Camera Parameters

NeRF--: Neural Radiance Fields Without Known Camera Parameters Project Page | Arxiv | Colab Notebook | Data Zirui Wang¹, Shangzhe Wu², Weidi Xie², Min

Active Vision Laboratory 411 Dec 26, 2022
Unofficial & improved implementation of NeRF--: Neural Radiance Fields Without Known Camera Parameters

[Unofficial code-base] NeRF--: Neural Radiance Fields Without Known Camera Parameters [ Project | Paper | Official code base ] ⬅️ Thanks the original

Jianfei Guo 239 Dec 22, 2022
NeRF Meta-Learning with PyTorch

NeRF Meta Learning With PyTorch nerf-meta is a PyTorch re-implementation of NeRF experiments from the paper "Learned Initializations for Optimizing Co

Sanowar Raihan 78 Dec 18, 2022
Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields.

This repository contains the code release for Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields. This implementation is written in JAX, and is a fork of Google's JaxNeRF implementation. Contact Jon Barron if you encounter any issues.

Google 625 Dec 30, 2022
This repository contains a PyTorch implementation of "AD-NeRF: Audio Driven Neural Radiance Fields for Talking Head Synthesis".

AD-NeRF: Audio Driven Neural Radiance Fields for Talking Head Synthesis | Project Page | Paper | PyTorch implementation for the paper "AD-NeRF: Audio

null 551 Dec 29, 2022
Code release for DS-NeRF (Depth-supervised Neural Radiance Fields)

Depth-supervised NeRF: Fewer Views and Faster Training for Free Project | Paper | YouTube Pytorch implementation of our method for learning neural rad

null 524 Jan 8, 2023
Putting NeRF on a Diet: Semantically Consistent Few-Shot View Synthesis Implementation

Putting NeRF on a Diet: Semantically Consistent Few-Shot View Synthesis Implementation This project attempted to implement the paper Putting NeRF on a

null 254 Dec 27, 2022
NeRF visualization library under construction

NeRF visualization library using PlenOctrees, under construction pip install nerfvis Docs will be at: https://nerfvis.readthedocs.org import nerfvis s

Alex Yu 196 Jan 4, 2023
A PyTorch implementation of NeRF (Neural Radiance Fields) that reproduces the results.

NeRF-pytorch NeRF (Neural Radiance Fields) is a method that achieves state-of-the-art results for synthesizing novel views of complex scenes. Here are

Yen-Chen Lin 3.2k Jan 8, 2023
D-NeRF: Neural Radiance Fields for Dynamic Scenes

D-NeRF: Neural Radiance Fields for Dynamic Scenes [Project] [Paper] D-NeRF is a method for synthesizing novel views, at an arbitrary point in time, of

Albert Pumarola 291 Jan 2, 2023
Code release for NeRF (Neural Radiance Fields)

NeRF: Neural Radiance Fields Project Page | Video | Paper | Data Tensorflow implementation of optimizing a neural representation for a single scene an

null 6.5k Jan 1, 2023
Pytorch implementation for A-NeRF: Articulated Neural Radiance Fields for Learning Human Shape, Appearance, and Pose

A-NeRF: Articulated Neural Radiance Fields for Learning Human Shape, Appearance, and Pose Paper | Website | Data A-NeRF: Articulated Neural Radiance F

Shih-Yang Su 172 Dec 22, 2022
A minimal TPU compatible Jax implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis

NeRF Minimal Jax implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis. Result of Tiny-NeRF RGB Depth

Soumik Rakshit 11 Jul 24, 2022
Build upon neural radiance fields to create a scene-specific implicit 3D semantic representation, Semantic-NeRF

Semantic-NeRF: Semantic Neural Radiance Fields Project Page | Video | Paper | Data In-Place Scene Labelling and Understanding with Implicit Scene Repr

Shuaifeng Zhi 243 Jan 7, 2023
Point-NeRF: Point-based Neural Radiance Fields

Point-NeRF: Point-based Neural Radiance Fields Project Sites | Paper | Primary c

Qiangeng Xu 662 Jan 1, 2023
[arXiv'22] Panoptic NeRF: 3D-to-2D Label Transfer for Panoptic Urban Scene Segmentation

Panoptic NeRF Project Page | Paper | Dataset Panoptic NeRF: 3D-to-2D Label Transfer for Panoptic Urban Scene Segmentation Xiao Fu*, Shangzhan zhang*,

Xiao Fu 111 Dec 16, 2022
Unofficial pytorch-lightning implement of Mip-NeRF

mipnerf_pl Unofficial pytorch-lightning implement of Mip-NeRF, Here are some results generated by this repository (pre-trained models are provided bel

Jianxin Huang 159 Dec 23, 2022
[arXiv'22] Panoptic NeRF: 3D-to-2D Label Transfer for Panoptic Urban Scene Segmentation

Panoptic NeRF: 3D-to-2D Label Transfer for Panoptic Urban Scene Segmentation Xiao Fu1*  Shangzhan Zhang1*  Tianrun Chen1  Yichong Lu1  Lanyun Zhu2  Xi

Xiao Fu 37 May 17, 2022