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^1
1MIT, ²Google, ³Nvidia
ICRA 2022

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.

Install the following dependencies before you run the code:

pip install -r requirements.txt

• COLMAP

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.

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

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

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.

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!

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

1. Launch the jupyter notebook:

jupyter notebook examine-correspondence.ipynb

2. 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.

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