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)