P2P-Bridge: Diffusion Bridges for 3D Point Cloud Denoising

Mathias Vogel¹
Keisuke Tateno², Marc Pollefeys^1,3, Federico Tombari^2,4, Marie-Julie Rakotosaona^*2 Francis Engelmann^*1,2,
¹ETH Zurich, ²Google, ³Microsoft, ⁴TUM,
^*Equal Contribution

ECCV 2024

📚Paper | 💾Code

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.

⚙️ Requirements

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

🗂️ Data Preparation

⚙️ Requirements

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

🧸 Object Datasets (PU-Net and PC-Net)

Download both zip files from ScoreDenoise. Extract them into data/objects such that the folder structure looks as follows:

data/objects
├── examples
├── PCNet
├── PUNet

🏠 Indoor Scene Datasets

To prepare the indoor scene datasets, follow the instructions here.

🚀 Training

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

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
└── ...

📊 Evaluation

🧸 PU-Net and PC-Net

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

🏠 Indoor Scenes

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.

1. Denoising:

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>

2. 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.

🧩 Denoise Your Own Data

🎥 Example

🏠 Real-World Data (Indoor Scenes)

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

🧸 Synthetic Data (Objects)

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>