init

README.md

@@ -1,8 +1,8 @@
 ---
-title: R2flow
-emoji: 🌖
-colorFrom: yellow
-colorTo: purple
+title: R2Flow
+emoji: 🚗
+colorFrom: indigo
+colorTo: green
 sdk: gradio
 sdk_version: 5.22.0
 app_file: app.py

app.py

@@ -0,0 +1,225 @@
+import re
+
+import einops
+import gradio as gr
+import matplotlib.cm as cm
+import numpy as np
+import plotly.graph_objects as go
+import torch
+import torch.nn.functional as F
+import torchdiffeq
+
+DESCRIPTION = """
+<div class="head">
+<div class="title">Fast LiDAR Data Generation with Rectified Flows</div>
+<div class="conference">ICRA 2025</div>
+<div class="authors">
+<a href="https://kazuto1011.github.io/" target="_blank" rel="noopener"> Kazuto Nakashima</a><sup>1</sup>
+&nbsp;&nbsp;&nbsp;
+<a> Xiaowen Liu</a><sup>1</sup>
+&nbsp;&nbsp;&nbsp;
+<a> Tomoya Miyawaki</a><sup>1</sup>
+&nbsp;&nbsp;&nbsp;
+<a> Yumi Iwashita</a><sup>2</sup>
+&nbsp;&nbsp;&nbsp;
+<a> Ryo Kurazume</a><sup>1</sup>
+</div>
+<div class="affiliations">
+<sup>1</sup>Kyushu University
+&nbsp;&nbsp;&nbsp;
+<sup>2</sup>NASA Jet Propulsion Laboratory
+</div>
+<div class="materials">
+<a href="https://kazuto1011.github.io/r2flow">Project</a> |
+<a href="https://arxiv.org/abs/2412.02241">Paper</a> |
+<a href="https://github.com/kazuto1011/r2flow">Code</a>
+</div>
+<br>
+<div class="description">
+This is a demo of our paper "Fast LiDAR Data Generation with Rectified Flows" accepted to ICRA 2025.<br>
+We propose <strong>R2Flow</strong>, a rectified flow-based LiDAR generative model which generate the LiDAR range/reflectance images.<br>
+</div>
+<br>
+</div>
+"""
+
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
+
+torch.set_grad_enabled(False)
+torch.backends.cudnn.benchmark = True
+device = torch.device(device)
+
+
+model_dict = {
+    "1-RF": "r2flow-kitti360-1rf",
+    "2-RF": "r2flow-kitti360-2rf",
+    "2-RF + 4-TD": "r2flow-kitti360-2rf-4td",
+    "2-RF + 2-TD": "r2flow-kitti360-2rf-2td",
+    "2-RF + 1-TD": "r2flow-kitti360-2rf-1td",
+}
+
+
+torch_hub_kwargs = dict(
+    repo_or_dir="kazuto1011/r2flow",
+    model="pretrained_r2flow",
+    device=device,
+    show_info=False,
+)
+
+
+def colorize(tensor: torch.Tensor, cmap_fn=cm.turbo):
+    colors = cmap_fn(np.linspace(0, 1, 256))[:, :3]
+    colors = torch.from_numpy(colors).to(tensor)
+    tensor = tensor.squeeze(1) if tensor.ndim == 4 else tensor
+    ids = (tensor * 256).clamp(0, 255).long()
+    tensor = F.embedding(ids, colors).permute(0, 3, 1, 2)
+    tensor = tensor.mul(255).clamp(0, 255).byte()
+    return tensor
+
+
+def model_verbose(model, nfe, progress):
+    handler = progress.tqdm(range(nfe), desc="Generating...")
+
+    def _model(t, x):
+        handler.update(1)
+        return model(t, x)
+
+    return _model
+
+
+def generate(nfe: int, solver: str, phase: str, progress=gr.Progress()):
+    model, lidar_utils, _ = torch.hub.load(config=model_dict[phase], **torch_hub_kwargs)
+
+    with torch.inference_mode():
+        x1 = torchdiffeq.odeint(
+            func=model_verbose(model, int(nfe), progress),
+            y0=torch.randn(1, model.in_channels, *model.resolution, device=device),
+            t=torch.linspace(0, 1, int(nfe) + 1, device=device),
+            method=solver,
+        )[-1]
+
+    depth = lidar_utils.restore_metric_depth(x1[:, [0]])
+    rflct = lidar_utils.denormalize(x1[:, [1]])
+    point = lidar_utils.convert_metric_depth(depth, format="cartesian")
+
+    z_min, z_max = -2, 0.5
+    z = (point[:, [2]] - z_min) / (z_max - z_min)
+    color = colorize(z.clamp(0, 1), cm.viridis) / 255
+    point = einops.rearrange(point, "1 c h w -> (h w) c").cpu().numpy()
+    color = einops.rearrange(color, "1 c h w -> (h w) c").cpu().numpy()
+    fig = go.Figure(
+        data=[
+            go.Scatter3d(
+                x=-point[..., 0],
+                y=-point[..., 1],
+                z=point[..., 2],
+                mode="markers",
+                marker=dict(size=1, color=color),
+            )
+        ],
+        layout=dict(
+            scene=dict(
+                xaxis=dict(showticklabels=False, visible=False),
+                yaxis=dict(showticklabels=False, visible=False),
+                zaxis=dict(showticklabels=False, visible=False),
+                aspectmode="data",
+            ),
+            margin=dict(l=0, r=0, b=0, t=0),
+            paper_bgcolor="white",
+            plot_bgcolor="white",
+        ),
+    )
+    depth = depth / lidar_utils.max_depth
+    depth = colorize(depth, cm.turbo)[0].permute(1, 2, 0).cpu().numpy()
+    rflct = colorize(rflct, cm.turbo)[0].permute(1, 2, 0).cpu().numpy()
+
+    model.cpu()
+    lidar_utils.cpu()
+    return depth, rflct, fig
+
+
+def setup_dropdown(value):
+    if "TD" in value:
+        solver_choices = ["euler"]
+        solver_default = "euler"
+        num_step = re.findall(r"(\d+)-TD", value)[0]
+        nfe_choices = [num_step]
+        nfe_default = num_step
+    else:
+        solver_choices = ["euler", "dopri5"]
+        solver_default = "euler"
+        nfe_choices = [2**i for i in range(0, 9)]
+        nfe_default = 256
+    dropdown_solver = gr.Dropdown(
+        choices=solver_choices,
+        value=solver_default,
+        label="ODE solver",
+        info="Fixed if TD enabled",
+    )
+    dropdown_nfe = gr.Dropdown(
+        choices=nfe_choices,
+        value=nfe_default,
+        label="Number of sampling steps",
+        info="Fixed if TD enabled",
+    )
+    return dropdown_solver, dropdown_nfe
+
+
+with gr.Blocks(
+    css="""
+.head {
+  text-align: center;
+  display: block;
+  font-size: var(--text-xl);
+}
+
+.title {
+  font-size: var(--text-xxl);
+  font-weight: bold;
+  margin-top: 2rem;
+}
+
+.description {
+  font-size: var(--text-lg);
+}
+    """,
+    theme=gr.themes.Ocean(),
+) as demo:
+    gr.HTML(DESCRIPTION)
+
+    with gr.Row(variant="panel"):
+        with gr.Column():
+            gr.Textbox(device, label="Running device")
+            dropdown_model = gr.Dropdown(
+                choices=list(model_dict.keys()),
+                value="2-RF + 4-TD",
+                label="Model checkpoint",
+                info="RF: rectified flow, TD: timestep distillation",
+            )
+            dropdown_solver, dropdown_nfe = setup_dropdown(dropdown_model.value)
+            dropdown_model.change(
+                setup_dropdown,
+                inputs=[dropdown_model],
+                outputs=[dropdown_solver, dropdown_nfe],
+            )
+            btn = gr.Button(value="Generate", variant="primary")
+
+        with gr.Column():
+            range_view = gr.Image(type="numpy", label="Range image")
+            rflct_view = gr.Image(type="numpy", label="Reflectance image")
+            point_view = gr.Plot(label="Point cloud")
+
+    btn.click(
+        generate,
+        inputs=[dropdown_nfe, dropdown_solver, dropdown_model],
+        outputs=[range_view, rflct_view, point_view],
+    )
+
+
+demo.queue()
+demo.launch()

pre-requirements.txt

@@ -0,0 +1,3 @@
+--index-url https://download.pytorch.org/whl/cpu
+torch==2.1.2
+torchvision==0.16.2

requirements.txt

@@ -0,0 +1,14 @@
+einops==0.6.1
+gradio==5.22.0
+kornia==0.7.0
+matplotlib==3.7.1
+pydantic==2.6.3
+rich==13.5.1
+simple-parsing==0.1.5
+torchcfm==1.0.5
+torchdiffeq==0.2.4
+tqdm==4.66.1
+plotly==6.0.1
+numpy==1.26.4
+--find-links https://shi-labs.com/natten/wheels
+natten==0.17.5+torch260cpu