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)