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