nerf_rcnn

NeRF-RCNN Training

We extend NeRF-RPN by adding a RoI head for predicting class labels as well as discrete 3D masks, for each region proposal generated by NeRF-RPN.

As a result, you need to first train a NeRF-RPN model following the instructions here, or use one of our pre-trained models if you are using Hypersim, 3D-FRONT, or ScanNet datasets.

Then, you can start to train the NeRF-RCNN RoI head using the NeRF-RPN model as the initial weights.

Data Preparation

To train a NeRF-RCNN, organize the dataset as below:

rcnn_data
|- masks
|  |- scene_00.npy
|  └-...
|- features
|  |- scene_00.npz
|  └-...
|- rois
|  |- scene_00.npz
|  └-...
|- metadata
|  |- scene_00.json
|  └-...
└- dataset_split.json

features should contain the same NeRF extracted RGB density features as used in NeRF-RPN. rois should contain the region proposals generated from a trained NeRF-RPN model with --save_results in the commandline.

metadata contains JSON files describing all instance information in each scene, including object bounding boxes, class labels, and object IDs. masks contains 3D masks of the same resolution as the feature grids, with object IDs as entries.

dataset_split.json contains all the scene names in training, validation, and test sets. Please refer to our released dataset for a example format of these files.

Training

We provide an example shell script train_rcnn.sh to run the training code. Some notes about the usage:

• Supply the path to the trained NeRF-RPN model weights to --rpn_ckpt. Make sure --rpn_type, --backbone_type, and all arguments starting with --rpn_* match with the setting you used to train the NeRF-RPN model.
• Change DATA_ROOT and the data split filename as needed.
• Currectly we only take RoIs in AABB format, which means the NeRF-RPN model also needs to be outputing AABB proposals.
• We suggest precompute the RoIs and supply them as input. It might be possible to generate region proposals on the fly during NeRF-RCNN training, but we haven't tested it.

Inference

To do inference on test set and get the instance labels and 3D masks, change the --mode to eval and supply the trained NeRF-RCNN model weights to --rcnn_ckpt. Run the same script again, the program will automatically save the instance 3D masks, boxes, classes, and scores for each scene in the test set as .npz files, to the path specified in --save_path. You can modify --save_top_k to change the number of instances to save.