NeRF-Supervision

Project Page | Video | Paper | Data

PyTorch implementation of NeRF-Supervision, an RGB-only, self-supervised pipeline for learning object-centric dense descriptors from neural radiance fields (NeRFs).

NeRF-Supervision: Learning Dense Object Descriptors from Neural Radiance Fields
Lin Yen-Chen¹,
Pete Florence²,
Jonathan T. Barron²,
Tsung-Yi Lin³,
Alberto Rodriguez¹,
Phillip Isola¹
¹MIT, ²Google, ³Nvidia

ICRA 2022

What This Repo Contains?

This repo hosts the code for step (a) Use RGB images to optimize NeRF and (b) Generate dense correspondences. For step (c) Train dense object descriptors, we refer users to pytorch-dense-correspondence.

Dependencies

Install the following dependencies before you run the code:

pip install -r requirements.txt

COLMAP

How to Run?

In the following, we will use fork as an example:

OBJECT=fork

To use other objects, one can simply set the variable OBJECT=your_favorite_object.

Download example data

cd data
gdown 16_y_Nnh19Qhml0bg9RYR-hav0YOpWKuw
unzip -q $OBJECT.zip
cd ..

Generate camera poses and sparse depth information using COLMAP

The data should look like the following:

├── data
│   ├── fork
│   │   ├── images
│   │   │   ├── 000001.png
│   │   │   ├── 000002.png
│   │   │   ├── ...

To generate camera poses and the sparse point cloud:

python imgs2poses.py --data_dir ./data/$OBJECT

Train NeRFs

To train a DS-NeRF:

python DS_NeRF/run_nerf.py --config DS_NeRF/configs/$OBJECT.txt

All the checkpoints and rendered images will be saved in ./logs/$OBJECT.

Generate The Correspondence Dataset

Once the training is done, run the following lines:

python DS_NeRF/run_nerf.py --config DS_NeRF/configs/$OBJECT.txt --render_only --render_train --render_factor 2

This should produce a new directory called logs/$OBJECT/renderonly_train_$ITER where $ITER corresponds to the number of training iterations. This directory contains the generated correspondence and will be used in the next step!

Examine The Correspondence Dataset

We provide the following interface to examine the correspondences generated by NeRFs.

Launch the jupyter notebook:

jupyter notebook examine-correspondence.ipynb

In the first cell, specify DATA_DIR with the directory we generated from the last step: logs/$OBJECT/renderonly_train_$ITER. For example, it can be logs/$OBJECT/renderonly_train_049999.

Citation

@inproceedings{yen2022nerfsupervision,
  title={{NeRF-Supervision}: Learning Dense Object Descriptors from Neural Radiance Fields},
  author={Lin Yen-Chen and Pete Florence and Jonathan T. Barron and Tsung-Yi Lin and Alberto Rodriguez and Phillip Isola},
  booktitle={IEEE Conference on Robotics and Automation ({ICRA})},
  year={2022}
}

This codebase is highly based on the following publication:

@article{kangle2021dsnerf,
  title={Depth-supervised NeRF: Fewer Views and Faster Training for Free},
  author={Deng, Kangle and Liu, Andrew and Zhu, Jun-Yan and Ramanan, Deva},
  journal={arXiv preprint arXiv:2107.02791},
  year={2021}
}