R2DM

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

LiDAR Data Synthesis with Denoising Diffusion Probabilistic Models
Kazuto Nakashima, Ryo Kurazume
ICRA 2024
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

• The initial run takes about 15 min to preprocess & cache the whole dataset.
• The default configuration is config H (R2DM) in our paper.
• Distributed training and mixed precision are enabled by default.
• Run with --help to list the available options.

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

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

• The discrete/continuous diffusion processes are based on lucidrains/denoising-diffusion-pytorch.
• The BEV-based evaluation metrics are based on vzyrianov/lidargen.