📢 NEWS: We have released MVEdit, an upgraded codebase based on SSDNeRF. MVEdit supports all SSDNeRF models and configs, and offers new features such as diffusers support and improved SSDNeRF GUI.

SSDNeRF

Official PyTorch implementation of the ICCV 2023 paper:

Single-Stage Diffusion NeRF: A Unified Approach to 3D Generation and Reconstruction <br> Hansheng Chen^1,*, Jiatao Gu², Anpei Chen³, Wei Tian¹, Zhuowen Tu⁴, Lingjie Liu⁵, Hao Su⁴<br> ¹Tongji University, ²Apple, ³ETH Zürich, ⁴UCSD, ⁵University of Pennsylvania <br> *Work done during a remote internship with UCSD.

[project page] [paper]

Part of this codebase is based on torch-ngp and MMGeneration. <br>

https://github.com/Lakonik/SSDNeRF/assets/53893837/22e7ee6c-7576-44f2-b408-41089180e359

Highlights

• Code to reproduce ALL the experiments in the paper and supplementary material (including single-view reconstruction on the real KITTI Cars dataset). <br><img src="ssdnerf_kitti.gif" width="500" alt=""/>
• New features including support for tiled triplanes (rollout layout), FP16 diffusion sampling, and 16-bit caching.
• A simple GUI demo (modified from torch-ngp). <br><img src="ssdnerf_gui.png" width="500" alt=""/>

Installation

Prerequisites

The code has been tested in the environment described as follows:

• Linux (tested on Ubuntu 18.04/20.04 LTS)
• Python 3.7
• CUDA Toolkit 11
• PyTorch 1.12.1
• MMCV 1.6.0
• MMGeneration 0.7.2

Also, this codebase should be able to work on Windows systems as well (tested in the inference mode).

Other dependencies can be installed via pip install -r requirements.txt.

An example of commands for installing the Python packages is shown below:

# Export the PATH of CUDA toolkit
export PATH=/usr/local/cuda-11.3/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH

# Create conda environment
conda create -y -n ssdnerf python=3.7
conda activate ssdnerf

# Install PyTorch (this script is for CUDA 11.3)
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch

# Install MMCV and MMGeneration
pip install -U openmim
mim install mmcv-full==1.6
git clone https://github.com/open-mmlab/mmgeneration && cd mmgeneration && git checkout v0.7.2
pip install -v -e .
cd ..

# Clone this repo and install other dependencies
git clone <this repo> && cd <repo folder>
pip install -r requirements.txt

Compile CUDA packages

There are two CUDA packages from torch-ngp that need to be built locally.

cd lib/ops/raymarching/
pip install -e .
cd ../shencoder/
pip install -e .
cd ../../..

Data preparation

Download srn_cars.zip and srn_chairs.zip from here. Unzip them to ./data/shapenet.

Download abo_tables.zip from here. Unzip it to ./data/abo. For convenience I have converted the ABO dataset into PixelNeRF's SRN format.

If you want to try single-view reconstruction on the real KITTI Cars dataset, please download the official KITTI 3D object dataset, including left color images, calibration files, training labels, and instance segmentations.

Extract the downloaded archives according to the following folder tree (or use symlinks).

./
├── configs/
├── data/
│   ├── shapenet/
│   │   ├── cars_test/
│   │   ├── cars_train/
│   │   ├── cars_val/
│   │   ├── chairs_test/
│   │   ├── chairs_train/
│   │   └── chairs_val/
│   ├── abo/
│   │   ├── tables_train/
│   │   └── tables_test/
│   └── kitti/
│       └── training/
│           ├── calib/
│           ├── image_2/
│           ├── label_2/
|           └── instance_2/
├── demo/
├── lib/
├── tools/
…

For FID and KID evaluation, run the following commands to extract the Inception features of the real images. (This script will use all the available GPUs on your machine, so remember to set CUDA_VISIBLE_DEVICES.)

CUDA_VISIBLE_DEVICES=0 python tools/inception_stat.py configs/paper_cfgs/ssdnerf_cars_uncond.py
CUDA_VISIBLE_DEVICES=0 python tools/inception_stat.py configs/paper_cfgs/ssdnerf_chairs_recons1v.py
CUDA_VISIBLE_DEVICES=0 python tools/inception_stat.py configs/paper_cfgs/ssdnerf_abotables_uncond.py

For KITTI Cars preprocessing, run the following command.

python tools/kitti_preproc.py

About the configs

Naming convention

ssdnerf_cars3v_uncond
   │      │      └── testing data: test unconditional generation
   │      └── training data: train on Cars dataset, using 3 views per scene
   └── training method: single-stage diffusion nerf training
  
stage2_cars_recons1v
   │     │      └── testing data: test 3D reconstruction from 1 view
   │     └── training data: train on Cars dataset, using all views per scene
   └── training method: stage 2 of two-stage training

Models in the main paper

| Config | Checkpoint | Iters | FID | LPIPS | Comments | |:---------------------------------------------------------------------------|:-----------------------------------------------------------------------------------------------:|:------:|:------------:|:-----:|:----------------------------------------------------------------------------------------------------------------------------------------------------------------| | ssdnerf_cars_uncond | gdrive | 1M | 11.08 ± 1.11 | - | | | ssdnerf_abotables_uncond | gdrive | 1M | 14.27 ± 0.66 | - | | | ssdnerf_cars_recons1v | gdrive | 80K | 16.39 | 0.078 | | | ssdnerf_chairs_recons1v | gdrive | 80K | 10.13 | 0.067 | | | ssdnerf_cars3v_uncond_1m | | 1M | | - | The first half of training before resetting the triplanes. | | ssdnerf_cars3v_uncond_2m | gdrive | 1M | 19.04 ± 1.10 | - | The second half of training after resetting the triplanes (requires training ssdnerf_cars3v_uncond_1m first). | | ssdnerf_cars3v_recons1v | | 80K | | 0.106 | | | stage1_cars_recons16v | | 400K | | | Ablation study, NeRF reconstruction stage. | | stage2_cars_uncond |