Home

Awesome

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 implemented 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

  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

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.

Additional notes

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}
}