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