Skip to content

play_log

Script to play back log file and generate animation.

Usage

python play_log.py --log_file filename.pkl

Log files in the required format is generated by render_evaluation.py.

main()

Entry point of the evaluation program.

Source code in sdfest/estimation/scripts/play_log.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def main() -> None:
    """Entry point of the evaluation program."""
    global reset
    global animation_queued
    parser = argparse.ArgumentParser(
        description="Play log file and generate animation."
    )
    parser.add_argument("--log_file", required=True)
    args = parser.parse_args()

    with open(args.log_file, "rb") as f:
        data = pickle.load(f)
        config = data["config"]
        log_entries = data["log"]

    pipeline = SDFPipeline(config)

    # precompute meshes for each timestep
    for log_entry in tqdm(log_entries):
        # print(log_entry)
        if "latent_shape" in log_entry:
            out_mesh = pipeline.generate_mesh(
                log_entry["latent_shape"], 1 / log_entry["scale_inv"], True
            )
            out_mesh.position = log_entry["position"][0].detach().cpu().numpy()
            out_mesh.orientation = log_entry["orientation"][0].detach().cpu().numpy()
            log_entry["mesh"] = out_mesh.get_transformed_o3d_geometry()
            log_entry["mesh"].paint_uniform_color([0.2, 0.4, 0.7])
            log_entry["pointcloud"] = log_entry["mesh"].sample_points_uniformly(20000)

    # visualize log entries
    cam_meshes = []
    pointclouds = []
    KEY_ESCAPE = 256
    vis = o3d.visualization.VisualizerWithKeyCallback()
    vis.register_key_callback(key=ord("A"), callback_func=reset_bounding_box)
    vis.register_key_callback(key=ord("R"), callback_func=toggle_realtime)
    vis.register_key_callback(key=ord("S"), callback_func=switch_reconstruction_type)
    vis.register_key_callback(key=ord("N"), callback_func=queue_animation)
    vis.register_key_callback(key=ord("C"), callback_func=toggle_color)
    vis.register_key_callback(key=ord("F"), callback_func=toggle_camera_frames)
    vis.register_key_callback(key=ord(" "), callback_func=toggle_pause)
    vis.register_key_callback(key=KEY_ESCAPE, callback_func=quit_program)
    vis.create_window(width=640, height=480)
    print(
        "Controls\n\ta: reset view point & bounding box\n"
        "\tr: toggle realtime\n"
        "\ts: switch reconstruction_type\n",
        "\tn: queue animation\n",
        "\tc: toggle color\n",
        "\tf: toggle camera frames\n",
        "\tspace: pause loop\n",
    )
    first = True
    while True:
        vis.clear_geometries()
        animation_folder = None
        animation_files = []
        animation_timestamps = []
        if animation_queued:
            animation_queued = False
            animation_folder = (
                f"animation_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"
            )
            os.makedirs(animation_folder, exist_ok=True)
        start_time = time.time()
        for log_entry in log_entries:
            vis.poll_events()
            vis.update_renderer()
            if realtime and animation_folder is None:
                while log_entry["timestamp"] > time.time() - start_time:
                    vis.poll_events()
                    vis.update_renderer()

            if "camera_positions" in log_entry:
                cam_meshes = []
                for t_c2w, quat_c2w in zip(
                    log_entry["camera_positions"], log_entry["camera_orientations"]
                ):
                    frame_mesh = synthetic.Mesh(
                        mesh=o3d.geometry.TriangleMesh.create_coordinate_frame(
                            size=0.1, origin=[0, 0, 0]
                        ),
                        rel_scale=True,
                    )
                    frame_mesh.position = t_c2w.cpu().numpy()
                    frame_mesh.orientation = quaternion_utils.quaternion_multiply(
                        quat_c2w.cpu(), torch.tensor([1.0, 0, 0, 0])
                    ).numpy()
                    cam_meshes.append(frame_mesh.get_transformed_o3d_geometry())

            if "depth_images" in log_entry:
                pointclouds = []
                for depth_image, color_image, t_c2w, quat_c2w in zip(
                    log_entry["depth_images"],
                    log_entry["color_images"],
                    log_entry["camera_positions"],
                    log_entry["camera_orientations"],
                ):
                    points_c = pointset_utils.depth_to_pointcloud(
                        depth_image, pipeline.cam, normalize=False
                    )
                    pointcloud_torch = (
                        quaternion_utils.quaternion_apply(quat_c2w, points_c) + t_c2w
                    )
                    pointcloud_numpy = pointcloud_torch.cpu().numpy()
                    pointcloud_o3d = o3d.geometry.PointCloud(
                        o3d.utility.Vector3dVector(pointcloud_numpy)
                    )
                    if color:
                        pointcloud_colors_torch = color_image[depth_image != 0]
                        pointcloud_colors_numpy = pointcloud_colors_torch.cpu().numpy()
                        pointcloud_o3d.colors = o3d.utility.Vector3dVector(
                            pointcloud_colors_numpy
                        )
                    else:
                        pointcloud_o3d.colors = o3d.utility.Vector3dVector(
                            np.ones_like(pointcloud_numpy) * np.array([1.0, 0.2, 0.2])
                        )

                    pointclouds.append(pointcloud_o3d)

            while True:
                geometries = [] + pointclouds
                if camera_frames:
                    geometries += cam_meshes
                if "mesh" in log_entry and reconstruction_type != "none":
                    geometries.append(log_entry[reconstruction_type])
                vis.clear_geometries()
                for i, geometry in enumerate(geometries):
                    vis.add_geometry(geometry, reset_bounding_box=first or reset)
                    if i == len(geometries) - 1:
                        reset = first = False
                if not pause:
                    break
                else:
                    vis.poll_events()
                    vis.update_renderer()

            if animation_folder is not None:
                vis.poll_events()
                vis.update_renderer()
                timestamp = log_entry["timestamp"]
                filename = f"{timestamp}.png"
                vis.capture_screen_image(os.path.join(animation_folder, filename))
                animation_files.append(filename)
                animation_timestamps.append(timestamp)

        # create constant framerate video
        if animation_folder is not None:
            fps = 30
            frame_folder = os.path.join(animation_folder, "constant_framerate")
            os.makedirs(frame_folder, exist_ok=True)
            video_name = f"{animation_folder}.mp4"
            current_frame = animation_files.pop(0)
            current_time = animation_timestamps.pop(0)
            frame_number = 0
            while animation_files:
                if animation_timestamps[0] <= current_time:
                    current_frame = animation_files.pop(0)
                    animation_timestamps.pop(0)
                copyfile(
                    os.path.join(animation_folder, current_frame),
                    os.path.join(frame_folder, f"{frame_number:06d}.png"),
                )

                current_time += 1 / fps
                frame_number += 1

            ffmpeg.input(
                os.path.join(frame_folder, "*.png"), pattern_type="glob", framerate=fps
            ).output(video_name).run()

queue_animation(_)

Switch between mesh and pointcloud reconstruction.

Source code in sdfest/estimation/scripts/play_log.py 
91
92
93
94
95
96
def queue_animation(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Switch between mesh and pointcloud reconstruction."""
    global animation_queued
    animation_queued = True
    print("Saving animation for next loop.")
    return False

quit_program(_)

Quit program.

Source code in sdfest/estimation/scripts/play_log.py 
36
37
38
def quit_program(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Quit program."""
    exit()

reset_bounding_box(_)

Schedule resetting of bounding box.

Source code in sdfest/estimation/scripts/play_log.py 
41
42
43
44
45
46
def reset_bounding_box(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Schedule resetting of bounding box."""
    global reset
    print("Reset bounding box.")
    reset = True
    return False

switch_reconstruction_type(_)

Switch between mesh and pointcloud reconstruction.

Source code in sdfest/estimation/scripts/play_log.py 
81
82
83
84
85
86
87
88
def switch_reconstruction_type(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Switch between mesh and pointcloud reconstruction."""
    global reconstruction_type
    cur_id = reconstruction_types.index(reconstruction_type)
    cur_id = (cur_id + 1) % len(reconstruction_types)
    reconstruction_type = reconstruction_types[cur_id]
    print(f"Reconstruction: {reconstruction_type}")
    return False

toggle_camera_frames(_)

Toggle camera frame visualization.

Source code in sdfest/estimation/scripts/play_log.py 
73
74
75
76
77
78
def toggle_camera_frames(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Toggle camera frame visualization."""
    global camera_frames
    camera_frames = not camera_frames
    print(f"Camera frames: {camera_frames}")
    return False

toggle_color(_)

Toggle color pointcloud.

Source code in sdfest/estimation/scripts/play_log.py 
65
66
67
68
69
70
def toggle_color(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Toggle color pointcloud."""
    global color
    color = not color
    print(f"Color: {color}")
    return False

toggle_pause(_)

Toggle pause.

Source code in sdfest/estimation/scripts/play_log.py 
57
58
59
60
61
62
def toggle_pause(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Toggle pause."""
    global pause
    pause = not pause
    print(f"Pause: {pause}")
    return False

toggle_realtime(_)

Toggle realtime playback.

Source code in sdfest/estimation/scripts/play_log.py 
49
50
51
52
53
54
def toggle_realtime(_: o3d.visualization.VisualizerWithKeyCallback) -> bool:
    """Toggle realtime playback."""
    global realtime
    realtime = not realtime
    print(f"Realtime: {realtime}")
    return False