Home

Awesome

R2DM

R2DM is a denoising diffusion probabilistic model (DDPM) for LiDAR range/reflectance generation based on the equirectangular representation.

sampled in 256 steps

LiDAR Data Synthesis with Denoising Diffusion Probabilistic Models<br> Kazuto Nakashima, Ryo Kurazume<br> ICRA 2024<br> project | arxiv | online demo

Quick demo

Python denpendencies:

pip install torch torchvision einops tqdm pydantic

Unconditional generation using the pre-trained model:

import torch

# Sampling
r2dm, lidar_utils, cfg = torch.hub.load("kazuto1011/r2dm", "pretrained_r2dm", device="cuda")
lidar_image = r2dm.sample(batch_size=1, num_steps=256)  # (batch size, 2, height, width)

# Postprocessing
lidar_image = lidar_utils.denormalize(lidar_image.clamp(-1, 1))  # [-1,1] -> [0,1]
range_image = lidar_utils.revert_depth(lidar_image[:, [0]])  # Range
rflct_image = lidar_image[:, [1]]  # Reflectance
point_cloud = lidar_utils.to_xyz(range_image)  # Point cloud

Setup

Python & CUDA

w/ conda framework:

conda env create -f environment.yaml
conda activate r2dm

If you are stuck with an endless installation, try libmamba for the conda solver.

Dataset

For training & evaluation, please download the KITTI-360 dataset (163 GB) and make a symlink:

ln -sf $PATH_TO_KITTI360_ROOT data/kitti_360/dataset

Please set the environment variable $HF_DATASETS_CACHE to cache the processed dataset (default: ~/.cache/huggingface/datasets).

Training

To start training DDPMs:

accelerate launch train.py

To monitor the training progress:

tensorboard --logdir logs/

To generate samples w/ a training checkpoint (*.pth) at $CHECKPOINT_PATH:

python generate.py --ckpt $CHECKPOINT_PATH

Evaluation

To generate, save, and evaluate samples:

accelerate launch sample_and_save.py --ckpt $CHECKPOINT_PATH --output_dir $OUTPUT_DIR
python evaluate.py --ckpt $CHECKPOINT_PATH --sample_dir $OUTPUT_DIR

The generated samples are saved in $OUTPUT_DIR.

Completion demo

python completion_demo.py --ckpt $CHECKPOINT_PATH

completion

Citation

If you find this code useful for your research, please cite our paper:

@inproceedings{nakashima2024lidar,
    title     = {LiDAR Data Synthesis with Denoising Diffusion Probabilistic Models},
    author    = {Kazuto Nakashima and Ryo Kurazume},
    year      = 2024,
    booktitle = {Proceedings of the International Conference on Robotics and Automation (ICRA)},
    pages     = {14724--14731}
}

Acknowledgements