update

app.py

@@ -1,16 +1,26 @@
+import einops
 import gradio as gr
 import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
+import plotly.graph_objects as go
 import torch
 import torch.nn.functional as F
-import einops
-import plotly.graph_objects as go
+
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
 
 torch.set_grad_enabled(False)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device(device)
+
 ddpm, lidar_utils, _ = torch.hub.load(
-    "kazuto1011/r2dm", "pretrained_r2dm", device=device
+    "kazuto1011/r2dm",
+    "pretrained_r2dm",
+    device=device,
 )
 
 
@@ -30,19 +40,7 @@ def render_point_cloud(output, cmap):
     rflct = output[:, [1]]
     point = lidar_utils.to_xyz(depth).cpu().numpy()
     point = einops.rearrange(point, "1 c h w -> c (h w)")
-    angle = lidar_utils.ray_angles.rad2deg()
-    label = [
-        f"depth: {float(d):.2f}m<br>"
-        + f"reflectance: {float(r):.2f}<br>"
-        + f"elevation: {float(e):.2f}°<br>"
-        + f"azimuth: {float(a):.2f}°"
-        for d, r, e, a in zip(
-            einops.rearrange(depth, "1 1 h w -> (h w)"),
-            einops.rearrange(rflct, "1 1 h w -> (h w)"),
-            einops.rearrange(angle[0, 0], "h w -> (h w)"),
-            einops.rearrange(angle[0, 1], "h w -> (h w)"),
-        )
-    ]
+    # angle = lidar_utils.ray_angles.rad2deg()
     fig = go.Figure(
         data=[
             go.Scatter3d(
@@ -59,8 +57,19 @@ def render_point_cloud(output, cmap):
                     cmin=-2,
                     cmax=0.5,
                 ),
-                text=label,
-                hoverinfo="text",
+                # text=[
+                #     f"depth: {float(d):.2f}m<br>"
+                #     + f"reflectance: {float(r):.2f}<br>"
+                #     + f"elevation: {float(e):.2f}°<br>"
+                #     + f"azimuth: {float(a):.2f}°"
+                #     for d, r, e, a in zip(
+                #         einops.rearrange(depth, "1 1 h w -> (h w)"),
+                #         einops.rearrange(rflct, "1 1 h w -> (h w)"),
+                #         einops.rearrange(angle[0, 0], "h w -> (h w)"),
+                #         einops.rearrange(angle[0, 1], "h w -> (h w)"),
+                #     )
+                # ],
+                # hoverinfo="text",
             )
         ],
         layout=dict(
@@ -88,7 +97,7 @@ def generate(num_steps, cmap_name, progress=gr.Progress()):
     for i in progress.tqdm(range(num_steps), desc="Generating LiDAR data"):
         step_t = steps[:, i]
         step_s = steps[:, i + 1]
-        x = ddpm.p_sample(x, step_t, step_s)
+        x = ddpm.p_step(x, step_t, step_s)
     return render_point_cloud(x, plt.colormaps.get_cmap(cmap_name))
 
 
@@ -96,14 +105,17 @@ with gr.Blocks() as demo:
     gr.Markdown(
         """
         # 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>
-        [[arXiv]](https://arxiv.org/abs/2309.09256) [[Code]](https://github.com/kazuto1011/r2dm)
+        ICRA 2024<br>
+        [[Project]](https://kazuto1011.github.io/r2dm/) [[arXiv]](https://arxiv.org/abs/2309.09256) [[Code]](https://github.com/kazuto1011/r2dm)
+
+        R2DM is a denoising diffusion probabilistic model (DDPM) for LiDAR range/reflectance generation based on the equirectangular representation.
         """
     )
     with gr.Row():
         with gr.Column():
+            gr.Textbox(device, label="Device")
             num_steps = gr.Dropdown(
                 choices=[2**i for i in range(2, 10)],
                 value=16,