add dynamic point cloud visulization

visualization_pcd.py

@@ -0,0 +1,166 @@
+"""Record3D visualizer
+"""
+import os
+import time
+from pathlib import Path
+from decord import VideoReader, cpu
+
+import numpy as np
+import tyro
+import viser
+import viser.extras
+import viser.transforms as tf
+from tqdm.auto import tqdm
+
+
+def main(
+    data_path: str = "/apdcephfs_cq10/share_1290939/vg_share/reynli/cache/video_demo/video_depth",
+    vid_name: str = "01_dog",
+    downsample_factor: int = 8,
+    max_frames: int = 100,
+    share: bool = False,
+    point_size = 0.01
+) -> None:
+    
+    server = viser.ViserServer()
+    if share:
+        server.request_share_url()
+
+    print("Loading frames!")
+    dis_path = data_path + '/' + vid_name + '.npz'
+    vid_path = data_path + '/' + vid_name + '_rgb.mp4'
+    # vid_path = data_path + '/' + vid_name + '.mp4'
+    
+    disp_map = np.load(dis_path)['depth'][:, :, :]
+    T = disp_map.shape[0]
+    H = disp_map.shape[1]
+    W = disp_map.shape[2]
+
+    disp_max = disp_map.max()
+    disp_min = disp_map.min()
+    disp_map = (disp_map - disp_min) / (disp_max - disp_min)
+
+    vr = VideoReader(vid_path, ctx=cpu(0))
+    vid = vr[:].asnumpy()[:, 0:H, 0:W]
+    fps = vr.get_avg_fps()
+    num_frames = min(max_frames, T)
+
+    # Add playback UI.
+    with server.gui.add_folder("Playback"):
+        gui_timestep = server.gui.add_slider(
+            "Timestep",
+            min=0,
+            max=num_frames - 1,
+            step=1,
+            initial_value=0,
+            disabled=True,
+        )
+        gui_next_frame = server.gui.add_button("Next Frame", disabled=True)
+        gui_prev_frame = server.gui.add_button("Prev Frame", disabled=True)
+        gui_playing = server.gui.add_checkbox("Playing", True)
+        gui_framerate = server.gui.add_slider(
+            "FPS", min=1, max=60, step=0.1, initial_value=fps
+        )
+        gui_framerate_options = server.gui.add_button_group(
+            "FPS options", ("10", "20", "30", "60")
+        )
+
+    # Frame step buttons.
+    @gui_next_frame.on_click
+    def _(_) -> None:
+        gui_timestep.value = (gui_timestep.value + 1) % num_frames
+
+    @gui_prev_frame.on_click
+    def _(_) -> None:
+        gui_timestep.value = (gui_timestep.value - 1) % num_frames
+
+    # Disable frame controls when we're playing.
+    @gui_playing.on_update
+    def _(_) -> None:
+        gui_timestep.disabled = gui_playing.value
+        gui_next_frame.disabled = gui_playing.value
+        gui_prev_frame.disabled = gui_playing.value
+
+    # Set the framerate when we click one of the options.
+    @gui_framerate_options.on_click
+    def _(_) -> None:
+        gui_framerate.value = int(gui_framerate_options.value)
+
+    prev_timestep = gui_timestep.value
+
+    # Toggle frame visibility when the timestep slider changes.
+    @gui_timestep.on_update
+    def _(_) -> None:
+        nonlocal prev_timestep
+        current_timestep = gui_timestep.value
+        with server.atomic():
+            frame_nodes[current_timestep].visible = True
+            frame_nodes[prev_timestep].visible = False
+        prev_timestep = current_timestep
+        server.flush()  # Optional!
+
+    # Load in frames.
+    server.scene.add_frame(
+        "/frames",
+        wxyz=tf.SO3.exp(np.array([0.0, 0.0, 0.0])).wxyz,
+        position=(0, 0, 0),
+        show_axes=False,
+    )
+    frame_nodes: list[viser.FrameHandle] = []
+    for i in tqdm(range(num_frames)):
+
+        # Add base frame.
+        frame_nodes.append(server.scene.add_frame(f"/frames/t{i}", show_axes=False))
+        
+        position_image = np.where(np.zeros([H, W]) == 0)
+        v = np.array(position_image[0]) 
+        u = np.array(position_image[1]) 
+        d = disp_map[i, v, u]
+
+        zc = 1.0 / (d + 0.1)
+        # zc = 1.0 / (d + 1e-8)
+    
+        xc = zc * (u - (W / 2.0)) / (W/2.)
+        yc = zc * (v - (H / 2.0)) / (H/2.)
+
+        zc -= 4 # disp_max * 0.2
+
+        points = np.stack((xc, yc, zc), axis=1)
+        colors = vid[i, v, u]
+
+        points = points[::downsample_factor]
+        colors = colors[::downsample_factor]
+
+        # Place the point cloud in the frame.
+        server.scene.add_point_cloud(
+            name=f"/frames/t{i}/point_cloud",
+            points=points,
+            colors=colors,
+            point_size=point_size,#0.007,
+            point_shape="rounded",
+        )
+
+    # Hide all but the current frame.
+    for i, frame_node in enumerate(frame_nodes):
+        frame_node.visible = i == gui_timestep.value
+
+    # Playback update loop.
+    prev_timestep = gui_timestep.value
+    while True:
+        if gui_playing.value:
+            gui_timestep.value = (gui_timestep.value + 1) % num_frames
+
+        time.sleep(1.0 / gui_framerate.value)
+
+
+if __name__ == "__main__":
+    tyro.cli(main(
+        # dir path of saved rgb.mp4 and disp.npz, modify it to your own dir
+        data_path="outputs/results_open_world/",
+        # sample name, modify it to your own sample name
+        vid_name="wukong",
+        # downsample factor of dense pcd
+        downsample_factor=8,
+        # point cloud size
+        point_size=0.007
+    ))