Mathias Vogel1
Keisuke Tateno2,
Marc Pollefeys1,3,
Federico Tombari2,4,
Marie-Julie Rakotosaona*2
Francis Engelmann*1,2,
1ETH Zurich,
2Google,
3Microsoft,
4TUM,
*Equal Contribution
P2P-Bridge introduces a novel approach for point cloud denoising by adapting Diffusion Schrödinger bridges to learn an optimal transport plan between paired point sets. Further enhancements are possible by incorporating additional features such as RGB data and point-wise DINOV2 features.
The code was tested using Python 3.10 and CUDA 11.8 on Ubuntu 22.04 and WSL2. Due to compatibility with older methods, there are quite a few dependencies, but we tried to make installation easier by providing a script and accumulating CUDA code as much as possible.
First, create a new environment (we use conda) and install the dependencies using the following commands:
conda create -n p2pb python=3.10
conda activate p2pb
We recommend to first install torch
and torchvision
using the following command:
conda install pytorch==2.1.2 torchvision==0.16.2 pytorch-cuda=11.8 -c pytorch -c nvidia --yes
followed by the installation of Pytorch3D and TorchCluster.
Finally, install all other dependencies and compile the custom CUDA code using the following command:
sh install.sh
For data preparation, additional libraries are used. The requirements can be installed using the following command from the data
directory:
pip install -r requirements_data.txt
Download both zip files from ScoreDenoise.
Extract them into data/objects
such that the folder structure looks as follows:
data/objects
├── examples
├── PCNet
├── PUNet
To prepare the indoor scene datasets, follow the instructions here.
We use wandb to track the training process. To use wandb run
wandb init
in the terminal to log into your account (you will be asked for your API key). If you want to disable it, just run
wandb disabled
before running the training script.
To train a model, adjust the config
file in the configs
directory according to your data directory and run the following command:
python train.py --config <CONFIG FILE> --save_dir <SAVE DIRECTORY> --wandb_project <WANDB PROJECT NAME> --wandb_entity <WANDB ENTITY NAME>
For all available arguments, run
python train.py --help
which will also show you how to train using multiple GPUs.
Pretrained models can be downloaded from here. Extract the files into the pretrained
directory such that the folder structure looks as follows:
pretrained
├── PVDL_ARK_XYZ/
│ ├── opt.yaml
│ └── step_100000.pth
└── ...
To run an evaluation on the PU-Net and PC-Net test data, run the following two commands to reproduce our paper results. The commands first run the denoising on the test data, followed by metrics calculation.
python evaluate_objects.py --model_path ./pretrained/PVDS_PUNet/latest.pth --dataset PUNet
python evaluate_objects.py --model_path ./pretrained/PVDS_PUNet/latest.pth --dataset PCNet
The outputs are stored in output_objects/<dataset>
together with the metrics. The output folder can be changed using the --output_root
argument. For all available arguments, run
python evaluate_objects.py --help
To reproduce results on the indoor scenes dataset, we provide the following instructions for ScanNet . The ARKitScenes dataset can be evaluated in the same way.
To denoise rooms from our ScanNet test set, you need to have the rooms specified in splits/snpp_test.txt
preprocessed (see here). For automatic evaluation, copy all snpp_test
scenes into a separate folder called snpp_evaluation
. Then you can use our script to denoise all test rooms:
sh scripts/denoise_snpp.sh <PATH TO snpp_evaluation>
To evaluate the denoised rooms, run the following command:
python evaluate_rooms.py --data_root <PATH TO snpp_evaluation> --dataset snpp
This will calculate the metrics for all prediction files and generate a csv
file in the predictions
folder inside snpp_evaluation
. Note that the commands above use our coordinate only model by default. If you want to evaluate the models using RGB and DINOV2 features, use the RGB or RGB_DINO checkpoints by providing the corresponding weigths in the --model_path
argument.
To denoise real-world data such as indoor scenes, you can use the following command:
python denoise_room.py --room_path <ROOM PATH> --model_path <MODEL PATH> --out_path <OUTPUT PATH>
If you want to use precalculated features use the --feature_name
argument. For all available arguments, run
python denoise_room.py --help
To denoise synthetic data, you can use the following command:
python denoise_object.py --data_path <PATH TO XYZ FILE> --save_path <OUTPUT FILE> --model_path <MODEL PATH>
@inproceedings{vogel2024p2pbridgediffusionbridges3d,
title={P2P-Bridge: Diffusion Bridges for 3D Point Cloud Denoising},
author={Mathias Vogel and Keisuke Tateno and Marc Pollefeys and Federico Tombari and Marie-Julie Rakotosaona and Francis Engelmann},
year={2024},
booktitle={European Conference on Computer Vision (ECCV)},
}