Record3D visualizer#

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

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