This is the official source code of the research paper Learning Online Multi-Sensor Depth Fusion.
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.
The code has been tested with Python 3.8.5
- Clone the repository and submodules to your local directory: In the following, we denote the local path to the SenFuNet codebase as "ROOT_FOLDER".
git clone --recursive https://github.com/tfy14esa/SenFuNet.git
- Create a python virtual environment:
python -m venv senfunet_env
- Activate the virtual environment:
source senfunet_env/bin/activate
- 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.
- 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.
- Install the dependencies:
pip install -r requirements.txt
- 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.
We provide three separate datasets that can be used with SenFuNet. The download links and instructions are provided below.
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.
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
ROOT_FOLDER/configs/fusion/replica.yaml
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
ROOT_FOLDER/configs/fusion/replica.yaml
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
ROOT_FOLDER/lists/desk.txt
PATH_TO_MVS_DEPTH PATH_TO_FOLDER_WITH_TOF_DEPTH PATH_TO_RGB_TIMESTAMP_FILE PATH_TO_TOF_TIMESTAMP_FILE CAMERA_MATRICES
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
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
PATH_TO_RGB PATH_TO_TOF_DEPTH PATH_TO_MVS_DEPTH CAMERA_MATRICES
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
setup_colmap.py
reconstruct_colmap_slurm_SCENE.sh
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:
- SETTINGS.experiment_path -> path where the logging is done and the models are saved
- DATA.root_dir -> Path to data folder
- DATA.train_scene_list -> specifies data used during training
- DATA.val_scene_list -> specifies data used during validation
- 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.
- TRAINING.routing_stereo_model_path -> path to psmnet stereo routing model (only used when training with a routing network)
- TRAINING.routing_tof_model_path -> path to tof routing model (only used when training with a routing network)
- 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.
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:
- SETTINGS.experiment_path -> path where the logging is done and the models are saved.
- DATA.root_dir -> Path to data folder
- DATA.train_scene_list -> specifies data used during training
- DATA.val_scene_list -> specifies data used during validation
- DATA.test_scene_list -> specifies data used during testing
- 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.
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
- ToF+PSMNet without routing. Use the config file located at ROOT_FOLDER/configs/fusion/replica.yaml
- 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
- 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]
- 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
- ToF+MVS. Use the config file located at ROOT_FOLDER/configs/fusion/corbs.yaml.
Scene3D
- ToF+MVS. Use the config file located at ROOT_FOLDER/configs/fusion/scene3d.yaml.
- 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]
In the event that you want to train your own routing network, use the following guide to train and test the routing network.
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:
- SETTINGS.experiment_path -> path where the logging is done and the models are saved.
- DATA.root_dir -> Path to data folder
- DATA.train_scene_list -> specifies data used during training
- DATA.val_scene_list -> specifies data used during validation
- DATA.test_scene_list -> specifies data used during 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
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.
In the event that you want to reproduce our baseline results on the Replica dataset, follow the steps outlined below.
The early fusion baseline is only applicable to the Replica dataset since it requires ground truth depth maps for training.
- Select the appropriate sensor suite i.e. ToF+PSMNet or SGM+PSMNet. Change the config variable appropriately.
DATA.input
- Specify the path to the routing network using the config variable For example, the early fusion routing network for ToF+PSMNet fusion is available at ROOT_FOLDER/models/routing/tof_psmnet/model/best.pth.tar
TESTING.routing_model_path
- Set
TESTING.use_outlier_filter: False FILTERING_MODEL.model: 'tsdf_early_fusion' ROUTING.do: True
- Test using the test_fusion.py script with the config as input.
- Select the appropriate sensor suite i.e. ToF+PSMNet or SGM+PSMNet. Change the config variable appropriately.
DATA.input
- Set
TESTING.use_outlier_filter: False
- Set
FILTERING_MODEL.model: 'tsdf_middle_fusion'
- Test using the test_fusion.py script with the config as input.
- Select the appropriate sensor suite i.e. ToF+PSMNet or SGM+PSMNet. Change the config variable appropriately.
DATA.input
- Specify the path to the fusion network using the config variable 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
TESTING.fusion_model_path
- Set
FUSION_MODEL.use_fusion_net: True FUSION_MODEL.extraction_strategy: 'trilinear_interpolation' TESTING.use_outlier_filter: False FILTERING_MODEL.model: 'routedfusion'
- 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 Note that this requires testing and saving the grid of this model first.
TESTING.routedfusion_nn_model:
- Test using the test_fusion.py script with the config as input.
We adapted some code from routedfusion. Thanks for making this code publicly available.
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 Erik Sandström for questions, comments and reporting bugs.