๐Ÿ˜ฎThe official implementation of "CoNeRF: Controllable Neural Radiance Fields" ๐Ÿ˜ฎ

Overview

CoNeRF: Controllable Neural Radiance Fields

arXiv MIT license Website Datasets

This is the official implementation for "CoNeRF: Controllable Neural Radiance Fields"

The codebase is based on HyperNeRF implemente in JAX, building on JaxNeRF.

Setup

The code can be run under any environment with Python 3.8 and above. (It may run with lower versions, but we have not tested it).

We recommend using Miniconda and setting up an environment:

conda create --name conerf python=3.8

Next, install the required packages:

pip install -r requirements.txt

Install the appropriate JAX distribution for your environment by following the instructions here. For example:

# For CUDA version 11.1
pip install --upgrade "jax[cuda111]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

Dataset

Basic structure

The dataset uses the same format as Nerfies for the image extraction and camera estimation.

For annotations, we create an additional file annotations.yml consisting of attribute values and their corresponding frames, and a folder with [frame_id].json files (only annotated frames are required to have a corresponding .json file) where each *.json file is a segmentation mask created with LabelMe. In summary, each dataset has to have the following structure:

<dataset>
    โ”œโ”€โ”€ annotations
    โ”‚   โ””โ”€โ”€ ${item_id}.json
    โ”œโ”€โ”€ annotations.yml
    โ”œโ”€โ”€ camera
    โ”‚   โ””โ”€โ”€ ${item_id}.json
    โ”œโ”€โ”€ camera-paths
    โ”œโ”€โ”€ colmap
    โ”œโ”€โ”€ rgb
    โ”‚   โ”œโ”€โ”€ ${scale}x
    โ”‚   โ””โ”€โ”€ โ””โ”€โ”€ ${item_id}.png
    โ”œโ”€โ”€ metadata.json
    โ”œโ”€โ”€ dataset.json
    โ”œโ”€โ”€ scene.json
    โ””โ”€โ”€ mapping.yml

The mapping.yml file can be created manually and serves to map class indices to class names which were created with LabelMe. It has the following format:

<index-from-0>: <class-name>

for example:

0: left eye
1: right eye

The annotations.yml can be created manually as well (though we encourage using the provided notebook for this task) and has the following format:

- class: <id>
  frame: <number>
  value: <attribute-value> # between -1 and 1

for example:

- class: 0 # corresponding to left eye
  frame: 128
  value: -1
- class: 1 # corresponding to right eye
  frame: 147
  value: 1
- class: 2 # corresponding to mouth
  frame: 147
  value: -1 

Principles of annotating the data

  • Our framework works well with just a bunch of annotations (for extreme points as an example). For our main face visualizations, we used just 2 annotations per attribute.
  • We highly recommend annotating these frames that are extremes of possible controllability, for example, fully eye closed will be -1 value and fully open eye will +1 value. Though it is not necessary to be exact in extremes, the more accurate annotations, the more accurate controllability you can expect
  • Each attribute can be annotated independently, i.e., there is no need to look for frames that have exactly extreme values of all attributes. For example, left eye=-1 and left eye=+1 values can be provided in frames 28 and 47, while right eye=-1 and right eye=+1 can be provided in any other frames.
  • Masks should be quite rough oversized, it is generally better to have bigger than smaller annotations.
  • The general annotation pipeline looks like this:
  1. Find set of frames that consist of extreme attributions (e.g. closed eye, open eye etc.).
  2. Provide necessary values in for attributes to be controlled in annotations.yml.
  3. Set names for these attributes (necessary for the masking part).
  4. Run LabelMe.
  5. Save annotated frames in annotations/.

Now you can run the training! Also, check out our datasets (52GB of data) to avoid any preprocessing steps on your own.

We tried our best to make our CoNeRF codebase to be general for novel view synthesis validation dataset (conerf/datasets/nerfies.py file) but we mainly focused on the interpolation task. If you have an access to the novel view synthesis rig as used in NeRFies or HyperNeRF, and you find out that something doesn't work, please leave an issue.

Providing value annotations

We extended the basic notebook used in NeRFies and HyperNeRF for processing the data so that you can annotate necessary images with attributes. Please check out notebooks/Capture_Processing.ipynb for more details. The notebook (despite all the files from NeRFies) will also generate <dataset>/annotations.yml and <dataset>/mapping.yml files.

Providing masking annotations

We adapted data loading class to handle annotations from LabelMe (we used its docker version). Example annotation for one of our datasets looks like this:

example-annotation

The program generates *.json files in File->Output Dir which should be located in <dataset>/annotations/ folder.

Training

After preparing a dataset, you can train a Nerfie by running:

export DATASET_PATH=/path/to/dataset
export EXPERIMENT_PATH=/path/to/save/experiment/to
python train.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_configs configs/test_local_attributes.gin

To plot telemetry to Tensorboard and render checkpoints on the fly, also launch an evaluation job by running:

python eval.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_configs configs/test_local_attributes.gin

The two jobs should use a mutually exclusive set of GPUs. This division allows the training job to run without having to stop for evaluation.

Configuration

  • We use Gin for configuration.
  • We provide a couple preset configurations.
  • Please refer to config.py for documentation on what each configuration does.
  • Preset configs:
    • baselines/: All configs that were used to perform quantitative evaluation in the experiments, including baseline methods. The _proj suffix denotes a method that uses a learnable projection.
      • ours.gin: The full CoNeRF architecture with masking.
      • hypernerf_ap[_proj].gin: The axis-aligned plane configuration for HyperNeRF.
      • hypernerf_ds[_proj].gin: The deformable surface configuration for HyperNeRF.
      • nerf_latent[_proj].gin: The configuration for a simple baselines where we concatenate a learnable latent with each coordinate (resembles HyperNeRF AP without the warping field).
      • nerfies[_proj].gin: The configuration for the NeRFies model.
      • nerf.gin: The configuration for the simplest NeRF architecture.
    • full-hd/, hd/ and post/: We repurposed our baselines/ours.gin configuration for training for different resolutions and different sampling parameters that increase the quality of the generated images. Using post/ours.gin required us to use 4x A100 GPU for 2 weeks to make the training converge.

Synthetic dataset

We generated the synthetic dataset using Kubric. You can find the generation script here. After generating the dataset, you can run prepare_kubric_dataset.py to canonicalize its format to the same one that works with CoNeRF. The dataset is already attached in the provided zip file.

Additional scripts

All scripts below are used as the ones for training, they need $EXPERIMENT_PATH and $DATASET_PATH to be specified. They save the results into $EXPERIMENT_PATH.

  • render_changing_attributes.py: Renders each of changing attributes under a fixed camera.
  • render_video.py: Renders changing view under a fixed set of attributes.
  • render_all.py: Renders dynamically changing attributes and the camera parameters.
  • train_lr.py: Estimates parameters of the linear regression. The estimated model maps highly dimensional embedding into controllable attributes.

Additional notes

  • We have used notebooks/Results.ipynb to generate tables/visualizations for the article. While it may not particularily useful for you case, we have left it so you can copy or reuse some of its snippets. It's especially useful because it shows how to extract data from tensorboards.
  • We removed some of notebooks that were available in the HyperNeRF's codebase (ex. for training) but were no longer applicable to CoNeRF. We highly recommend using available scripts. If you have ever managed to adapt HyperNeRF's notebooks, please leave a pull request.

Citing

If you find our work useful, please consider citing:

@inproceedings{kania2022conerf,
  title     = {{CoNeRF: Controllable Neural Radiance Fields}},
  author    = {Kania, Kacper and Yi, Kwang Moo and Kowalski, Marek and Trzci{\'n}ski, Tomasz and Tagliasacchi, Andrea},
  booktitle   = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  year      = {2022}
}
Owner
Kacper Kania
PhDing in Neural Human Rendering ... ๐Ÿ‘€
Kacper Kania
Semantic segmentation task for ADE20k & cityscapse dataset, based on several models.

semantic-segmentation-tensorflow This is a Tensorflow implementation of semantic segmentation models on MIT ADE20K scene parsing dataset and Cityscape

HsuanKung Yang 83 Oct 13, 2022
Implementation of the final project of the course DDA6309 Probabilistic Graphical Model

Task-aware Joint CWS and POS (TCwsPos) This is the implementation of the final project of the course DDA6309 Probabilistic Graphical Models, The Chine

Peng 1 Dec 26, 2021
PyTorch implementation of Progressive Growing of GANs for Improved Quality, Stability, and Variation.

PyTorch implementation of Progressive Growing of GANs for Improved Quality, Stability, and Variation. Warning: the master branch might collapse. To ob

559 Dec 14, 2022
A code repository associated with the paper A Benchmark for Rough Sketch Cleanup by Chuan Yan, David Vanderhaeghe, and Yotam Gingold from SIGGRAPH Asia 2020.

A Benchmark for Rough Sketch Cleanup This is the code repository associated with the paper A Benchmark for Rough Sketch Cleanup by Chuan Yan, David Va

33 Dec 18, 2022
Unofficial implementation of MUSIQ (Multi-Scale Image Quality Transformer)

MUSIQ: Multi-Scale Image Quality Transformer Unofficial pytorch implementation of the paper "MUSIQ: Multi-Scale Image Quality Transformer" (paper link

41 Jan 02, 2023
AAAI 2022: Stationary diffusion state neural estimation

Stationary Diffusion State Neural Estimation Although many graph-based clustering methods attempt to model the stationary diffusion state in their obj

็ปฝ็จ 33 Nov 24, 2022
InterfaceGAN++: Exploring the limits of InterfaceGAN

InterfaceGAN++: Exploring the limits of InterfaceGAN Authors: Apavou Clรฉment & Belkada Younes From left to right - Images generated using styleGAN and

Younes Belkada 42 Dec 23, 2022
CS_Final_Metal_surface_detection - This is a final project for CoderSchool Machine Learning bootcamp on 29/12/2021.

CS_Final_Metal_surface_detection This is a final project for CoderSchool Machine Learning bootcamp on 29/12/2021. The project is based on the dataset

Cuong Vo 1 Dec 29, 2021
Paddle pit - Rethinking Spatial Dimensions of Vision Transformers

ๅŸบไบŽPaddleๅฎž็ŽฐPiT โ€”โ€”Rethinking Spatial Dimensions of Vision Transformers,arxiv ๅฎ˜ๆ–นๅŽŸ็‰ˆไปฃ

Hongtao Wen 4 Jan 15, 2022
Codes for "CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation"

CSDI This is the github repository for the NeurIPS 2021 paper "CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation

106 Jan 04, 2023
VSR-Transformer - This paper proposes a new Transformer for video super-resolution (called VSR-Transformer).

VSR-Transformer By Jiezhang Cao, Yawei Li, Kai Zhang, Luc Van Gool This paper proposes a new Transformer for video super-resolution (called VSR-Transf

Jiezhang Cao 225 Nov 13, 2022
ProFuzzBench - A Benchmark for Stateful Protocol Fuzzing

ProFuzzBench - A Benchmark for Stateful Protocol Fuzzing ProFuzzBench is a benchmark for stateful fuzzing of network protocols. It includes a suite of

155 Jan 08, 2023
Text to image synthesis using thought vectors

Text To Image Synthesis Using Thought Vectors This is an experimental tensorflow implementation of synthesizing images from captions using Skip Though

Paarth Neekhara 2.1k Jan 05, 2023
Notebooks, slides and dataset of the CorrelAid Machine Learning Winter School

CorrelAid Machine Learning Winter School Welcome to the CorrelAid ML Winter School! Task The problem we want to solve is to classify trees in Roosevel

CorrelAid 12 Nov 23, 2022
Implementations of paper Controlling Directions Orthogonal to a Classifier

Classifier Orthogonalization Implementations of paper Controlling Directions Orthogonal to a Classifier , ICLR 2022, Yilun Xu, Hao He, Tianxiao Shen,

Yilun Xu 33 Dec 01, 2022
A Python library for generating new text from existing samples.

ReMarkov is a Python library for generating text from existing samples using Markov chains. You can use it to customize all sorts of writing from birt

8 May 17, 2022
Official release of MSHT: Multi-stage Hybrid Transformer for the ROSE Image Analysis of Pancreatic Cancer axriv: http://arxiv.org/abs/2112.13513

MSHT: Multi-stage Hybrid Transformer for the ROSE Image Analysis This is the official page of the MSHT with its experimental script and records. We de

Tianyi Zhang 53 Dec 27, 2022
A particular navigation route using satellite feed and can help in toll operations & traffic managemen

How about adding some info that can quanitfy the stress on a particular navigation route using satellite feed and can help in toll operations & traffic management The current analysis is on the satel

Ashish Pandey 1 Feb 14, 2022
Free like Freedom

This is all very much a work in progress! More to come! ( We're working on it though! Stay tuned!) Installation Open an Anaconda Prompt (in Windows, o

2.3k Jan 04, 2023
Tutorial on active learning with the Nvidia Transfer Learning Toolkit (TLT).

Active Learning with the Nvidia TLT Tutorial on active learning with the Nvidia Transfer Learning Toolkit (TLT). In this tutorial, we will show you ho

Lightly 25 Dec 03, 2022