Skip to content

Code for Learning Online Multi-Sensor Depth Fusion - SenFuNet available here: https://arxiv.org/abs/2204.03353

Notifications You must be signed in to change notification settings

eriksandstroem/SenFuNet

Repository files navigation

SenFuNet

This is the official source code of the research paper Learning Online Multi-Sensor Depth Fusion.

Architecture

Many hand-held or mixed reality devices typically use a single sensor for 3D reconstruction although they often comprise multiple sensors. Multi-sensor depth fusion is able to substantially improve the robustness and accuracy of 3D reconstruction methods, but existing techniques are not robust enough to handle sensors which operate in different domains. To this end, we introduce SenFuNet, a depth fusion approach that learns sensor-specific noise and outlier statistics and combines the data streams of depth frames from different sensors in an online fashion. Our method fuses multi-sensor depth streams regardless of synchronization and calibration and generalizes well with little training data.

Please check out this video which describes our method and shows the most important results.

This repository provides all necessary steps to reproduce the results in the paper including data preparation and trained models that can be tested by the user directly.

Prerequisites

The code has been tested with Python 3.8.5

Installation

  1. Clone the repository and submodules to your local directory: In the following, we denote the local path to the SenFuNet codebase as "ROOT_FOLDER".
  2. Create a python virtual environment:
    python -m venv senfunet_env
  3. Activate the virtual environment:
    source senfunet_env/bin/activate
  4. Open the requirements.txt and change the path to the submodule evaluate-3d-reconstruction according to your file structure. Set the variable "ground_truth_data_base" to the directory where you store the ground truth .ply meshes (see below under the header "Ground Truth Meshes"). This variable is located in the config.py file of the evaluate_3d_reconstruction library located in the "deps" folder.
  5. Open the requirements.txt and replace the Open3D entry with a version of your choice. Install version 0.13.0 or newer (this has not been tested). Replace the entry with e.g. "open3d==0.13.0". A local version of the Open3D library was used when developing the project which is the reason why this entry needs to be replaced.
  6. Install the dependencies:
    pip install -r requirements.txt
  7. Note that the project uses "Weights and Biases" for logging during training. If you want to train your own models, therefore create your own account at Weights and Biases.

Data Preparation

We provide three separate datasets that can be used with SenFuNet. The download links and instructions are provided below.

Ground Truth Meshes

Download the ground truth meshes used for F-score evaluation here.

Note: be sure to set the variable "ground_truth_data_base" to the directory where you store the ground truth .ply meshes. This variable is located in the config.py file of the evaluate_3d_reconstruction library located in the "deps" folder. This step has to be performed before installing the module via pip.

While not needed for this codebase, there is also the option to set the path to the tranformation folder where transformation matrices are stored which aligns the ground truth mesh and the predicted mesh before F-score evaluation.

Replica

Train Dataset: room 0, room 2, office 3, office 1, apartment 1, frl apartment 0

Validation Dataset: frl apartment 1

Test Dataset: office 0, hotel 0, office 4

Important: Store your dataset at a location specified in the config variable

DATA.root_dir
of the
ROOT_FOLDER/configs/fusion/replica.yaml
config file.

Quick Setup: In the event that you don't want to donwload the full dataset at first, only download the office 0 scene and use this for training, validation and testing. In that case, please change the paths of the variables

DATA.train_scene_list, DATA.val_scene_list, DATA.test_scene_list
listed in the config file
ROOT_FOLDER/configs/fusion/replica.yaml
.

CoRBS

The CoRBS dataset can be downloaded here. The CoRBS dataset does not include multi-view stereo (MVS) depth maps nor the ground truth signed distance grids (SDF). We provide these separately.

Train and Validation Dataset: Download the D1 trajectory of the desk scene. The corresponding MVS depth maps and ground truth SDF grid are available here.

Test Dataset: Download the H1 trajectory of the human scene. The corresponding MVS depth maps and ground truth SDF grid are available here.

Prepare the data such that the paths to the depth sensors and camera matrices are provided in the files

ROOT_FOLDER/lists/human.txt
and
ROOT_FOLDER/lists/desk.txt
For each row the entries are separated by spaces and structured as follows:
PATH_TO_MVS_DEPTH PATH_TO_FOLDER_WITH_TOF_DEPTH PATH_TO_RGB_TIMESTAMP_FILE PATH_TO_TOF_TIMESTAMP_FILE CAMERA_MATRICES

Scene3D

Download the stonewall and copy room scenes of the Scene3D dataset available here. Next, use the script

ROOT_FOLDER/data/save_every_tenth_frame.py
to save every tenth sample (we only integrate every tenth frame).

Next, download the MVS depth sensor for both scenes and ground truth SDF grid for the stonewall training scene. We were not able to construct a good ground truth SDF grid for the copy room scene thus only the F-score evaluation is accurate at test time. Download links: stonewall, copy room.

Arrange the data such that the paths listed in the corresponding

ROOT_FOLDER/lists/*.txt
file match your folder structure. For each for, the entries are separated by spaces and structured as follows:
PATH_TO_RGB PATH_TO_TOF_DEPTH PATH_TO_MVS_DEPTH CAMERA_MATRICES
.

Generate Multi-View Stereo Depth

In the event that you want to reproduce or generate your own MVS depth sensors, we provide the scripts for this. These are available in the folder

ROOT_FOLDER/data/mvs_depth_estimation
. First use the script
setup_colmap.py
and then the script
reconstruct_colmap_slurm_SCENE.sh
to use to generate the MVS depth maps. For information, we refer to the colmap documentation.

Training

To train SenFuNet, execute the script:

python train_fusion.py --ROOT_FOLDER/configs/fusion/CONFIG.yaml

where CONFIG is a config file.

The following paths need to be specified in the CONFIG.yaml file:

  1. SETTINGS.experiment_path -> path where the logging is done and the models are saved
  2. DATA.root_dir -> Path to data folder
  3. DATA.train_scene_list -> specifies data used during training
  4. DATA.val_scene_list -> specifies data used during validation
  5. DATA.test_scene_list -> specifies data used during testing

For the Replica dataset, in addition, specify the following paths. Note: only the Replica dataset can be trained and tested with depth denoising. For the remainder of this guide, we refer to the denoising network as the routing network which is the term used in the original paper RoutedFusion.

  1. TRAINING.routing_stereo_model_path -> path to psmnet stereo routing model (only used when training with a routing network)
  2. TRAINING.routing_tof_model_path -> path to tof routing model (only used when training with a routing network)
  3. TRAINING.routing_tof_2_model_path -> same path as "routing_tof_model_path". Used when doing multi-agent reconstruction.

The routing models are located at ROOT_FOLDER/models/routing.

Testing

To test SenFuNet, execute the script:

python test_fusion.py --ROOT_FOLDER/configs/fusion/CONFIG.yaml

where CONFIG is a config file.

The following paths need to be specified in the CONFIG.yaml file:

  1. SETTINGS.experiment_path -> path where the logging is done and the models are saved.
  2. DATA.root_dir -> Path to data folder
  3. DATA.train_scene_list -> specifies data used during training
  4. DATA.val_scene_list -> specifies data used during validation
  5. DATA.test_scene_list -> specifies data used during testing
  6. TESTING.fusion_model_path -> path to model to be tested. If routing was used during training, this model will include the routing network parameters as well.

The fusion models are located at ROOT_FOLDER/models/fusion.

Configs to Reproduce Results in Paper

Note: The F-scores reported in the paper are computed using meshes which are produced using the marching cubes implementation of Open3D. Gaining access to this required some changes of the C++ source code and to make installation simpler, we resort to the skimage implementation. Only minor numerical differences exist between the two implementations.

Replica

  1. ToF+PSMNet without routing. Use the config file located at ROOT_FOLDER/configs/fusion/replica.yaml
  2. ToF+PSMNet with routing. Modify the config file ROOT_FOLDER/configs/fusion/replica.yaml as follows:
ROUTING.do: True
TESTING.fusion_model_path: ROOT_FOLDER/models/fusion/tof_psmnet_routing/model/best.pth.tar
  1. SGM+PSMNet without routing. Modify the config file ROOT_FOLDER/configs/fusion/replica.yaml as follows:
TESTING.fusion_model_path: ROOT_FOLDER/models/fusion/sgm_psmnet/model/best.pth.tar
DATA.input: [sgm_stereo, stereo]
  1. SGM+PSMNet with routing. Modify the config file ROOT_FOLDER/configs/fusion/replica.yaml as follows:
ROUTING.do: True
TESTING.fusion_model_path: ROOT_FOLDER/models/fusion/sgm_psmnet_routing/model/best.pth.tar
DATA.input: [sgm_stereo, stereo]

CoRBS

  1. ToF+MVS. Use the config file located at ROOT_FOLDER/configs/fusion/corbs.yaml.

Scene3D

  1. ToF+MVS. Use the config file located at ROOT_FOLDER/configs/fusion/scene3d.yaml.
  2. ToF+ToF (multi-agent reconstruction). Modify the config file ROOT_FOLDER/configs/fusion/scene3d.yaml as follows:
TESTING.fusion_model_path: ROOT_FOLDER/models/fusion/tof_tof_scene3d_collab_rec/model/best.pth.tar
DATA.collaborative_reconstruction: True
DATA.input: [tof, tof_2]

Routing Network

In the event that you want to train your own routing network, use the following guide to train and test the routing network.

Training

Execute the following command to train a routing network:

python train_routing.py --ROOT_FOLDER/configs/routing/replica.yaml

where ROOT_FOLDER is the path to the SenFuNet folder.

The following paths need to be specified in the replica.yaml file:

  1. SETTINGS.experiment_path -> path where the logging is done and the models are saved.
  2. DATA.root_dir -> Path to data folder
  3. DATA.train_scene_list -> specifies data used during training
  4. DATA.val_scene_list -> specifies data used during validation
  5. DATA.test_scene_list -> specifies data used during testing

Testing

Execute the following command to test your routing network:

python test_routing.py --ROOT_FOLDER/configs/routing/replica.yaml

where ROOT_FOLDER is the path to the SenFuNet folder.

In addition to specifying the paths for training, this requires the model path variable

TESTING.fusion_model_path
to be specified in the replica.yaml config file.

The test script creates the output 16-bit depth images from the routing network but does not evaluate the depth accuracy. Quantitative evaluation of the depth maps needs to be done as a post processing step.

Baseline Methods

In the event that you want to reproduce our baseline results on the Replica dataset, follow the steps outlined below.

Early Fusion

The early fusion baseline is only applicable to the Replica dataset since it requires ground truth depth maps for training.

  1. Select the appropriate sensor suite i.e. ToF+PSMNet or SGM+PSMNet. Change the config variable
    DATA.input
    appropriately.
  2. Specify the path to the routing network using the config variable
    TESTING.routing_model_path
    For example, the early fusion routing network for ToF+PSMNet fusion is available at ROOT_FOLDER/models/routing/tof_psmnet/model/best.pth.tar
  3. Set
    TESTING.use_outlier_filter: False
    FILTERING_MODEL.model: 'tsdf_early_fusion'
    ROUTING.do: True
  4. Test using the test_fusion.py script with the config as input.

TSDF Fusion

  1. Select the appropriate sensor suite i.e. ToF+PSMNet or SGM+PSMNet. Change the config variable
    DATA.input
    appropriately.
  2. Set
    TESTING.use_outlier_filter: False
  3. Set
    FILTERING_MODEL.model: 'tsdf_middle_fusion'
  4. Test using the test_fusion.py script with the config as input.

RoutedFusion

  1. Select the appropriate sensor suite i.e. ToF+PSMNet or SGM+PSMNet. Change the config variable
    DATA.input
    appropriately.
  2. Specify the path to the fusion network using the config variable
    TESTING.fusion_model_path
    For example, the fusion network for ToF+PSMNet fusion without routing is available at ROOT_FOLDER/models/fusion/tof_psmnet_routedfusion/model/best.pth.tar
  3. Set
    FUSION_MODEL.use_fusion_net: True
    FUSION_MODEL.extraction_strategy: 'trilinear_interpolation'
    TESTING.use_outlier_filter: False
    FILTERING_MODEL.model: 'routedfusion'
  4. Note that to avoid excessive outliers from the RoutedFusion model, to produce the mesh from the predicted TSDF grid, we apply the nearest neighbor mask from the nearest neighbor extraction strategy and not the trilinear interpolation strategy. This requires loading a separate weight grid from a model of the same sensor suite but using nearest neighbor extraction. Set the name of the config variable
    TESTING.routedfusion_nn_model:
    Note that this requires testing and saving the grid of this model first.
  5. Test using the test_fusion.py script with the config as input.

Acknowledgement

We adapted some code from routedfusion. Thanks for making this code publicly available.

Citation

If you find our code or paper useful, please cite

@inproceedings{sandstrom2022learning,
  title={Learning Online Multi-Sensor Depth Fusion},
  author={Sandstr{\"o}m, Erik and Oswald, Martin R and Kumar, Suryansh and Weder, Silvan and Yu, Fisher and Sminchisescu, Cristian and Van Gool, Luc},
  booktitle = {Proceedings of the European Conference on Computer Vision (ECCV)},
  year={2022}
}

Contact

Contact Erik Sandström for questions, comments and reporting bugs.

About

Code for Learning Online Multi-Sensor Depth Fusion - SenFuNet available here: https://arxiv.org/abs/2204.03353

Resources

Stars

Watchers

Forks