Update app.py

app.py

@@ -579,35 +579,58 @@ def preload_current_trajectory(preset_name, custom_repo_id, custom_filename, tra
     return status
 
 
-def _compose_video_frame(gallery_items, frame_label):
+def _compose_video_frame(gallery_items, frame_label, status_plot=None):
+    """Compose one video frame.
+
+    Top: selected observation images.
+    Bottom: trajectory-status plot with the moving timestep cursor.
+    """
     if not gallery_items:
-        canvas = Image.new("RGB", (640, 360), color=(20, 20, 20))
-        draw = ImageDraw.Draw(canvas)
+        obs_canvas = Image.new("RGB", (640, 360), color=(20, 20, 20))
+        draw = ImageDraw.Draw(obs_canvas)
         draw.text((16, 16), "No selected image keys", fill=(255, 255, 255))
-        return np.asarray(canvas)
-
-    pil_images = []
-    for img, label in gallery_items:
-        pil_img = Image.fromarray(np.asarray(img, dtype=np.uint8)).convert("RGB")
-        label_h = 24
-        panel = Image.new("RGB", (pil_img.width, pil_img.height + label_h), color=(0, 0, 0))
-        panel.paste(pil_img, (0, label_h))
-        draw = ImageDraw.Draw(panel)
-        draw.text((6, 4), str(label), fill=(255, 255, 255))
-        pil_images.append(panel)
-
-    gap = 8
-    top_h = 28
-    width = sum(im.width for im in pil_images) + gap * max(len(pil_images) - 1, 0)
-    height = max(im.height for im in pil_images) + top_h
-    canvas = Image.new("RGB", (width, height), color=(0, 0, 0))
-    draw = ImageDraw.Draw(canvas)
-    draw.text((8, 6), frame_label, fill=(255, 255, 255))
-
-    x = 0
-    for im in pil_images:
-        canvas.paste(im, (x, top_h))
-        x += im.width + gap
+    else:
+        pil_images = []
+        for img, label in gallery_items:
+            pil_img = Image.fromarray(np.asarray(img, dtype=np.uint8)).convert("RGB")
+            label_h = 24
+            panel = Image.new("RGB", (pil_img.width, pil_img.height + label_h), color=(0, 0, 0))
+            panel.paste(pil_img, (0, label_h))
+            draw = ImageDraw.Draw(panel)
+            draw.text((6, 4), str(label), fill=(255, 255, 255))
+            pil_images.append(panel)
+
+        gap = 8
+        top_h = 28
+        width = sum(im.width for im in pil_images) + gap * max(len(pil_images) - 1, 0)
+        height = max(im.height for im in pil_images) + top_h
+        obs_canvas = Image.new("RGB", (width, height), color=(0, 0, 0))
+        draw = ImageDraw.Draw(obs_canvas)
+        draw.text((8, 6), frame_label, fill=(255, 255, 255))
+
+        x = 0
+        for im in pil_images:
+            obs_canvas.paste(im, (x, top_h))
+            x += im.width + gap
+
+    if status_plot is not None:
+        status_img = Image.fromarray(np.asarray(status_plot, dtype=np.uint8)).convert("RGB")
+
+        target_w = obs_canvas.width
+        if status_img.width != target_w:
+            target_h = max(1, int(round(status_img.height * float(target_w) / float(status_img.width))))
+            status_img = status_img.resize((target_w, target_h), resample=Image.Resampling.BILINEAR)
+
+        gap_h = 8
+        canvas = Image.new(
+            "RGB",
+            (max(obs_canvas.width, status_img.width), obs_canvas.height + gap_h + status_img.height),
+            color=(0, 0, 0),
+        )
+        canvas.paste(obs_canvas, (0, 0))
+        canvas.paste(status_img, (0, obs_canvas.height + gap_h))
+    else:
+        canvas = obs_canvas
 
     pad_w = int(np.ceil(canvas.width / 16.0) * 16)
     pad_h = int(np.ceil(canvas.height / 16.0) * 16)
@@ -619,7 +642,7 @@ def _compose_video_frame(gallery_items, frame_label):
     return np.asarray(canvas)
 
 
-def build_current_trajectory_video(preset_name, custom_repo_id, custom_filename, traj_id, image_keys, display_scale, reverse_channels, fps):
+def build_current_trajectory_video(preset_name, custom_repo_id, custom_filename, traj_id, image_keys, display_scale, reverse_channels, fps, valid_window_len):
     if imageio is None:
         return None, "Video export requires imageio and imageio-ffmpeg in requirements.txt."
 
@@ -659,14 +682,22 @@ def build_current_trajectory_video(preset_name, custom_repo_id, custom_filename,
                 bool(reverse_channels),
             )
             label = "trajectory {} | frame {}/{}".format(traj_id, t, len(traj) - 1)
-            frame = _compose_video_frame(gallery_items, label)
+            status_plot, _is_valid_start, _num_valid_starts = get_cached_status_plot(
+                repo_id,
+                filename,
+                traj_id,
+                t,
+                int(valid_window_len),
+            )
+            frame = _compose_video_frame(gallery_items, label, status_plot=status_plot)
             writer.append_data(frame)
     finally:
         writer.close()
 
-    status = "Built trajectory video"
+    status = "Built trajectory video with moving status plot"
     status += "\nTrajectory: {}".format(traj_id)
     status += "\nFrames: {} | FPS: {}".format(len(traj), fps)
+    status += "\nValid-window length: {}".format(int(valid_window_len))
     return out_path, status
 
 
@@ -853,15 +884,37 @@ def build_app():
             video_fps = gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Video FPS")
 
         preload_status = gr.Textbox(label="Preload / video status", lines=4, value="Not preloaded yet.", interactive=False)
-        trajectory_video = gr.Video(label="Trajectory video: smooth browser-side playback")
-
-        gallery = gr.Gallery(label="Observation images", columns=2, height="auto", object_fit="contain")
-        status_plot = gr.Image(label="Trajectory status: no_teacher_action and valid starts", type="numpy")
-        info = gr.Textbox(label="Frame info", lines=16)
 
-        with gr.Accordion("Debug: HDF5 tree", open=False):
-            inspect_btn = gr.Button("Inspect HDF5 structure")
-            hdf5_tree = gr.Textbox(lines=24, label="HDF5 tree")
+        with gr.Tabs():
+            with gr.Tab("Frame viewer"):
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        gallery = gr.Gallery(
+                            label="Observation images",
+                            columns=2,
+                            height="auto",
+                            object_fit="contain",
+                        )
+                    with gr.Column(scale=2):
+                        status_plot = gr.Image(
+                            label="Trajectory status: no_teacher_action and valid starts",
+                            type="numpy",
+                        )
+                info = gr.Textbox(label="Frame info", lines=12)
+
+            with gr.Tab("Trajectory video"):
+                trajectory_video = gr.Video(
+                    label="Trajectory video: observation images + moving status plot",
+                    height=520,
+                )
+                gr.Markdown(
+                    "The video is generated from the selected trajectory and image keys. "
+                    "Use this tab for smooth browser-side playback/scrubbing."
+                )
+
+            with gr.Tab("Debug"):
+                inspect_btn = gr.Button("Inspect HDF5 structure")
+                hdf5_tree = gr.Textbox(lines=24, label="HDF5 tree")
 
         preset.change(
             fn=update_custom_visibility,
@@ -925,7 +978,7 @@ def build_app():
 
         video_btn.click(
             fn=build_current_trajectory_video,
-            inputs=[preset, custom_repo_id, custom_filename, traj_slider, image_keys, display_scale, reverse_channels, video_fps],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, image_keys, display_scale, reverse_channels, video_fps, chunk_len],
             outputs=[trajectory_video, preload_status],
         )