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)