Record3D visualizerΒΆ

Parse and stream record3d captures. To get the demo data, see ./assets/download_record3d_dance.sh.

  1import time
  2from pathlib import Path
  3
  4import numpy as np
  5import tyro
  6from tqdm.auto import tqdm
  7
  8import viser
  9import viser.extras
 10import viser.transforms as tf
 11
 12
 13def main(
 14    data_path: Path = Path(__file__).parent / "assets/record3d_dance",
 15    downsample_factor: int = 4,
 16    max_frames: int = 100,
 17    share: bool = False,
 18) -> None:
 19    server = viser.ViserServer()
 20    if share:
 21        server.request_share_url()
 22
 23    print("Loading frames!")
 24    loader = viser.extras.Record3dLoader(data_path)
 25    num_frames = min(max_frames, loader.num_frames())
 26
 27    # Add playback UI.
 28    with server.gui.add_folder("Playback"):
 29        gui_point_size = server.gui.add_slider(
 30            "Point size",
 31            min=0.001,
 32            max=0.02,
 33            step=1e-3,
 34            initial_value=0.01,
 35        )
 36        gui_timestep = server.gui.add_slider(
 37            "Timestep",
 38            min=0,
 39            max=num_frames - 1,
 40            step=1,
 41            initial_value=0,
 42            disabled=True,
 43        )
 44        gui_next_frame = server.gui.add_button("Next Frame", disabled=True)
 45        gui_prev_frame = server.gui.add_button("Prev Frame", disabled=True)
 46        gui_playing = server.gui.add_checkbox("Playing", True)
 47        gui_framerate = server.gui.add_slider(
 48            "FPS", min=1, max=60, step=0.1, initial_value=loader.fps
 49        )
 50        gui_framerate_options = server.gui.add_button_group(
 51            "FPS options", ("10", "20", "30", "60")
 52        )
 53
 54    # Frame step buttons.
 55    @gui_next_frame.on_click
 56    def _(_) -> None:
 57        gui_timestep.value = (gui_timestep.value + 1) % num_frames
 58
 59    @gui_prev_frame.on_click
 60    def _(_) -> None:
 61        gui_timestep.value = (gui_timestep.value - 1) % num_frames
 62
 63    # Disable frame controls when we're playing.
 64    @gui_playing.on_update
 65    def _(_) -> None:
 66        gui_timestep.disabled = gui_playing.value
 67        gui_next_frame.disabled = gui_playing.value
 68        gui_prev_frame.disabled = gui_playing.value
 69
 70    # Set the framerate when we click one of the options.
 71    @gui_framerate_options.on_click
 72    def _(_) -> None:
 73        gui_framerate.value = int(gui_framerate_options.value)
 74
 75    prev_timestep = gui_timestep.value
 76
 77    # Toggle frame visibility when the timestep slider changes.
 78    @gui_timestep.on_update
 79    def _(_) -> None:
 80        nonlocal prev_timestep
 81        current_timestep = gui_timestep.value
 82        with server.atomic():
 83            # Update point size.
 84            point_nodes[current_timestep].point_size = gui_point_size.value
 85
 86            # Toggle visibility.
 87            frame_nodes[current_timestep].visible = True
 88            frame_nodes[prev_timestep].visible = False
 89        prev_timestep = current_timestep
 90        server.flush()  # Optional!
 91
 92    # Load in frames.
 93    server.scene.add_frame(
 94        "/frames",
 95        wxyz=tf.SO3.exp(np.array([np.pi / 2.0, 0.0, 0.0])).wxyz,
 96        position=(0, 0, 0),
 97        show_axes=False,
 98    )
 99    frame_nodes: list[viser.FrameHandle] = []
100    point_nodes: list[viser.PointCloudHandle] = []
101    for i in tqdm(range(num_frames)):
102        frame = loader.get_frame(i)
103        position, color = frame.get_point_cloud(downsample_factor)
104
105        # Add base frame.
106        frame_nodes.append(server.scene.add_frame(f"/frames/t{i}", show_axes=False))
107
108        # Place the point cloud in the frame.
109        point_nodes.append(
110            server.scene.add_point_cloud(
111                name=f"/frames/t{i}/point_cloud",
112                points=position,
113                colors=color,
114                point_size=gui_point_size.value,
115                point_shape="rounded",
116            )
117        )
118
119        # Place the frustum.
120        fov = 2 * np.arctan2(frame.rgb.shape[0] / 2, frame.K[0, 0])
121        aspect = frame.rgb.shape[1] / frame.rgb.shape[0]
122        server.scene.add_camera_frustum(
123            f"/frames/t{i}/frustum",
124            fov=fov,
125            aspect=aspect,
126            scale=0.15,
127            image=frame.rgb[::downsample_factor, ::downsample_factor],
128            wxyz=tf.SO3.from_matrix(frame.T_world_camera[:3, :3]).wxyz,
129            position=frame.T_world_camera[:3, 3],
130        )
131
132        # Add some axes.
133        server.scene.add_frame(
134            f"/frames/t{i}/frustum/axes",
135            axes_length=0.05,
136            axes_radius=0.005,
137        )
138
139    # Hide all but the current frame.
140    for i, frame_node in enumerate(frame_nodes):
141        frame_node.visible = i == gui_timestep.value
142
143    # Playback update loop.
144    prev_timestep = gui_timestep.value
145    while True:
146        if gui_playing.value:
147            gui_timestep.value = (gui_timestep.value + 1) % num_frames
148
149        time.sleep(1.0 / gui_framerate.value)
150
151
152if __name__ == "__main__":
153    tyro.cli(main)