• Stars
    star
    101
  • Rank 338,166 (Top 7 %)
  • Language
    Python
  • Created about 1 year ago
  • Updated about 1 year ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Canonical Factors for Hybrid Neural Fields @ ICCV 2023

TILTED

Project page β€’ arXiv

Code release for our ICCV 2023 paper:

Brent Yi1, Weijia Zeng1, Sam Buchanan2, and Yi Ma1. Canonical Factors for Hybrid Neural Fields. International Conference on Computer Vision (ICCV), 2023.
1UC Berkeley, 2TTI-Chicago

Overview

We study neural field architectures that rely on factored feature volumes, by (1) analyzing factored grids in 2D to characterize undesirable biases for axis-aligned signals, and (2) using the resulting insights to study TILTED, a family of hybrid neural field architectures that removes these biases.

This repository is structured as follows:

.
β”œβ”€β”€ tilted
β”‚   β”œβ”€β”€ core                - Code shared between experiments. Factored grid
β”‚   β”‚                         and neural decoder implementations.
β”‚   β”œβ”€β”€ nerf                - Neural radiance field rendering, training, and
β”‚   β”‚                         dataloading utilities.
β”‚   β”œβ”€β”€ rgb2d               - 2D image reconstruction data and training
β”‚   β”‚                         utilities.
β”‚   └── sdf                 - Signed distance field dataloading, training, and
β”‚                             meshing infrastructure.
β”‚
β”œβ”€β”€ paper_commands          - Commands used for running paper experiments (NeRF)
β”œβ”€β”€ paper_results           - Output files used to generate paper tables. (NeRF)
β”‚                             Contains hyperparameters, evaluation metrics,
β”‚                             runtimes, etc.
β”‚
β”œβ”€β”€ tables_nerf.ipynb       - Table generation notebook for NeRF experiments.
β”‚
β”œβ”€β”€ train_nerf.py           - Training script for neural radiance field experiments.
β”œβ”€β”€ visualize_nerf.py       - Visualize trained neural radiance fields.
β”‚
└── requirements.txt        - Python dependencies.

Note that training scripts for 2D and SDF experiments have not yet been released. Feel free to reach out if you need these.

Running

Setup

This repository has been tested with Python 3.8, jax==0.4.9, and jaxlib==0.4.9+cuda11.cudnn86. We recommend first installing JAX via their official instructions: https://github.com/google/jax#installation

We've packaged dependencies into a requirements.txt file:

pip install -r requirements.txt

Visualization

We use Tensorboard for logging.

After training, radiance fields can be interactively visualized. Helptext for the visualization script can be found via:

python visualize_nerf.py --help

As a runnable example, we've uploaded trained checkpoints for the Kitchen dataset here.

This can be unzipped in tilted/ and visualized via:

# Checkpoints can be selected via the dropdown on the right.
# The 'Reset Up Direction' button will also be when orbitting / panning!
python visualize_nerf.py ./example_checkpoints

The visualization script supports RGB, PCA, and feature norm visualization:

TILTED.Visualizer.mov

The core viewer infrastructure has been moved into nerfstudio-project/viser, which may be helpful if you're interested in visualization for other projects.

Datasets

Meshes for SDF experiments were downloaded from alecjacobson/common-3d-test-models/.

All NeRF datasets were downloaded using nerfstudio's ns-download-data command:

# Requires nerfstudio installation.
ns-download-data blender
ns-download-data nerfstudio all

Training

Commands we used for training NeRF models in the paper can be found in paper_commands/.

Here are two examples, which should run at ~65 it/sec on an RTX 4090:

# Train a model on a synthetic scene.
python train_nerf.py blender-kplane-32c-axis-aligned --dataset-path {path_to_data}

# Train a model on a real scene.
python train_nerf.py nerfstudio-kplane-32c-axis-aligned --dataset-path {path_to_data}

The --help flag can also be passed in to print helptext.

Notes

This is research code, so parts of it may be chaotic. We've put effort into refactor and cleanup before release, but there's always more work to do here! If you have questions or comments, please reach out.

Some notes:

  • The global orientation can have a large impact on performance of baselines. --render-config.global-rotate-seed INT can be set in train_nerf.py to try a different global orientation; paper results sweep across 0, 1, and 2 for each synthetic scene.
  • For speeding things up, the bottleneck training step count can be dropped significantly without hurting performance. This is dictated by --bottleneck.optim.projection-decay-start and --bottleneck.optim.projection-decay-steps; bottleneck training stop as soon as the projection LR hits 0.
  • Runtimes can vary significantly between machines. Our experiments were run using JAX 0.4.9 and CUDA 11.8 on RTX 4090 GPUs.

This material is based upon work supported by the National Science Foundation Graduate Research Fellowship Program under Grant DGE 2146752. YM acknowledges partial support from the ONR grant N00014-22-1-2102, the joint Simons Foundation-NSF DMS grant 2031899, and a research grant from TBSI.

If any of it is useful, you can also cite:

@inproceedings{tilted2023,
    author = {Yi, Brent and Zeng, Weijia and Buchanan, Sam and Ma, Yi},
    title = {Canonical Factors for Hybrid Neural Fields},
    booktitle = {International Conference on Computer Vision (ICCV)},
    year = {2023},
}

More Repositories

1

tyro

CLI interfaces & config objects, from types
Python
462
star
2

jaxlie

Rigid transforms + Lie groups in JAX
Python
220
star
3

jaxls

Sparse nonlinear least squares for JAX
Python
168
star
4

dfgo

Differentiable Factor Graph Optimization for Learning Smoothers @ IROS 2021
Python
78
star
5

jax_dataclasses

Pytrees + dataclasses ❀️
Python
59
star
6

tensorf-jax

Unofficial implementation of Tensorial Radiance Fields (Chen & Xu β€˜22)
Python
38
star
7

pips-jax

JAX port of Persistent Independent Particles
Python
36
star
8

multimodalfilter

Jupyter Notebook
27
star
9

jelly_mechanical

Solidworks files for ME135 quadruped project
21
star
10

jax-ekf

Generic EKF, with support for non-Euclidean manifolds
Python
20
star
11

isort.vim

Async isort plugin for Vim + Neovim
Vim Script
20
star
12

minGPT-flax

GPT implementation in Flax
Python
18
star
13

stl_web_viewer2

Javascript utility for embedding 3D models
JavaScript
12
star
14

sparky_firmware

Firmware for RHex-style robot
C++
9
star
15

brushless_driver

BLDC motor driver design; 48V, 3.5A continuous / 6.5A peak
HTML
7
star
16

dotfilesp

Configuration files & setup scripts
Vim Script
6
star
17

fannypack

Tools for training PyTorch models
Python
5
star
18

jax_cuda_boilerplate

Toy package for custom CUDA kernels + JAX
Python
5
star
19

fifteen

Python
4
star
20

as5047d_breakout

Absolute encoder breakout PCB for prototyping
HTML
3
star
21

drawing_machine_firmware

Firmware + G-code parser for plotter
C
3
star
22

jelly2_mechanical

2
star
23

as5048b_breakout

daisy chainable absolute magnetic encoder
HTML
2
star
24

stl_web_viewer

deprecated in favor of https://github.com/brentyi/stl_web_viewer2
JavaScript
2
star
25

jaxfg

new version: https://github.com/brentyi/jaxls
Python
2
star
26

crossmodal_filtering1

obsolete; see https://github.com/stanford-iprl-lab/torchfilter
Jupyter Notebook
1
star
27

digikey_parser

nodejs package for scraping part information from digikey barcodes
JavaScript
1
star
28

marshmello_web

HTML
1
star
29

opthex

Python
1
star
30

170_proj

Python
1
star
31

keyboard

Mechanical keyboard PCB designs
1
star
32

jax-ldr

(unofficial) Closed-Loop Data Transcription to an LDR via Minimaxing Rate Reduction
Jupyter Notebook
1
star