RetrievalFuse

Paper | Project Page | Video

RetrievalFuse: Neural 3D Scene Reconstruction with a Database
Yawar Siddiqui, Justus Thies, Fangchang Ma, Qi Shan, Matthias Nießner, Angela Dai
ICCV2021

This repository contains the code for the ICCV 2021 paper RetrievalFuse, a novel approach for 3D reconstruction from low resolution distance field grids and from point clouds.

In contrast to traditional generative learned models which encode the full generative process into a neural network and can struggle with maintaining local details at the scene level, we introduce a new method that directly leverages scene geometry from the training database.

File and Folders

Broad code structure is as follows:

Further, the data/ directory has the following layout

data                    # root data directory
├── sdf_008             # low-res (8^3) distance fields
    ├── <dataset_0>     
        ├── <sample_0>
        ├── <sample_1>
        ├── <sample_2>
        ...
    ├── <dataset_1>
    ...
├── sdf_016             # low-res (16^3) distance fields
    ├── <dataset_0>
        ├── <sample_0>
        ├── <sample_1>
        ├── <sample_2>
        ...
    ├── <dataset_1>
    ...
├── sdf_064             # high-res (64^3) distance fields
    ├── <dataset_0>
            ├── <sample_0>
            ├── <sample_1>
            ├── <sample_2>
            ...
        ├── <dataset_1>
        ...
├── pc_20K              # point cloud inputs
    ├── <dataset_0>
        ├── <sample_0>
        ├── <sample_1>
        ├── <sample_2>
        ...
    ├── <dataset_1>
    ...
├── splits              # train/val splits
├── size                # data needed by SceneHandler class (autocreated on first run)
├── occupancy           # data needed by SceneHandler class (autocreated on first run)

Dependencies

pip install -r requirements.txt

Be sure that you pull the ChamferDistancePytorch submodule in external.

Data Preparation

For ShapeNetV2 and Matterport, get the appropriate meshes from the datasets. For 3DFRONT get the 3DFUTURE meshes and 3DFRONT scripts. For getting 3DFRONT meshes use our fork of 3D-FRONT-ToolBox to create room meshes.

Once you have the meshes, use our fork of sdf-gen to create distance field low-res inputs and high-res targets. For creating point cloud inputs simply use trimesh.sample.sample_surface (check util/misc/sample_scene_point_clouds). Place the processed data in appropriate directories:

• data/sdf_008/<dataset> or data/sdf_016/<dataset> for low-res inputs

• data/pc_20K/<dataset> for point clouds inputs

• data/sdf_064/<dataset> for targets

Training the Retrieval Network

Make sure that CUDA_HOME variable is set. To train retrieval networks use the following command:

python trainer/train_retrieval.py --config config/<config> --val_check_interval 5 --experiment retrieval --wandb_main --sanity_steps 1

We provide some sample configurations for retrieval.

For super-resolution, e.g.

• config/super_resolution/ShapeNetV2/retrieval_008_064.yaml
• config/super_resolution/3DFront/retrieval_008_064.yaml
• config/super_resolution/Matterport3D/retrieval_016_064.yaml

For surface-reconstruction, e.g.

• config/surface_reconstruction/ShapeNetV2/retrieval_128_064.yaml
• config/surface_reconstruction/3DFront/retrieval_128_064.yaml
• config/surface_reconstruction/Matterport3D/retrieval_128_064.yaml

Once trained, create the retrievals for train/validation set using the following commands:

python util/retrieval.py  --mode map --retrieval_ckpt <trained_retrieval_ckpt> --config <retrieval_config>

python util/retrieval.py --mode compose --retrieval_ckpt <trained_retrieval_ckpt> --config <retrieval_config> 

Training the Refinement Network

Use the following command to train the refinement network

python trainer/train_refinement.py --config <config> --val_check_interval 5 --experiment refinement --sanity_steps 1 --wandb_main --retrieval_ckpt <retrieval_ckpt>

Again, sample configurations for refinement are provided in the config directory.

For super-resolution, e.g.

• config/super_resolution/ShapeNetV2/refinement_008_064.yaml
• config/super_resolution/3DFront/refinement_008_064.yaml
• config/super_resolution/Matterport3D/refinement_016_064.yaml

For surface-reconstruction, e.g.

• config/surface_reconstruction/ShapeNetV2/refinement_128_064.yaml
• config/surface_reconstruction/3DFront/refinement_128_064.yaml
• config/surface_reconstruction/Matterport3D/refinement_128_064.yaml

Visualizations and Logs

Visualizations and checkpoints are dumped in the runs/<experiment> directory. Logs are uploaded to the user's Weights&Biases dashboard.

Processed Data & Models (ShapeNet)

Download processed data for ShapeNetV2 dataset using the following command

bash data/download_shapenet_processed.sh

This will populate the data/sdf_008, data/sdf_064, data/pc_20K, data/occupancy and data/size folders with processed ShapeNet data.

To download trained models on ShapeNetV2 use the following script

bash data/download_shapenet_models.sh

This downloads the checkpoints for retrieval and refinement for ShapeNet on both super-resolution and surface reconstruction tasks, plus the already computed retrievals. You can resume training these with the --resume flag in appropriate scripts (or inference with --sanity_steps flag). E.g. for resuming (and / or dumping inferences from data/splits/ShapeNetV2/main/val_vis.txt) use the following command

# super-resolution
python trainer/train_refinement.py --config config/super_resolution/ShapeNetV2/refinement_008_064.yaml  --sanity_steps -1 --resume runs/checkpoints/superres_refinement_ShapeNetV2.ckpt --retrieval_ckpt runs/07101959_superresolution_ShapeNetV2_upload/_ckpt_epoch=79.ckpt --current_phase 3 --max_epoch 161 --new_exp_for_resume

# surface-reconstruction
python trainer/train_refinement.py --config config/surface_reconstruction/ShapeNetV2/refinement_128_064.yaml  --sanity_steps -1 --resume runs/checkpoints/surfacerecon_refinement_ShapeNetV2.ckpt --retrieval_ckpt runs/07101959_surface_reconstruction_ShapeNetV2_upload/_ckpt_epoch=59.ckp --current_phase 3 --max_epoch 161 --new_exp_for_resume

Citation

@inproceedings{siddiqui2021retrievalfuse,
  title = {RetrievalFuse: Neural 3D Scene Reconstruction with a Database},
  author = {Siddiqui, Yawar and Thies, Justus and Ma, Fangchang and Shan, Qi and Nie{\ss}ner, Matthias and Dai, Angela},
  booktitle = {Proc. International Conference on Computer Vision (ICCV)},
  month = oct,
  year = {2021},
  doi = {},
  month_numeric = {10}
}

License