Update app.py

app.py

@@ -487,7 +487,7 @@ def _make_trajectory_status_plot(traj, timestep, min_seq_len):
     no_teacher, valid_indices = _compute_valid_start_indices(traj, min_seq_len)
     is_valid_start = timestep in set(valid_indices)
 
-    fig, ax = plt.subplots(figsize=(9, 2.4), dpi=140)
+    fig, ax = plt.subplots(figsize=(10.5, 2.8), dpi=170)
 
     ax.step(
         np.arange(total_steps),
@@ -510,8 +510,8 @@ def _make_trajectory_status_plot(traj, timestep, min_seq_len):
     ax.axvline(timestep, color="black", linestyle="-", alpha=0.85, linewidth=1.5)
     ax.set_xlim(0, max(total_steps - 1, 1))
     ax.set_ylim(-0.38, 1.1)
-    ax.set_ylabel("Flag")
-    ax.set_xlabel("Timestep index")
+    ax.set_ylabel("Flag", fontsize=10)
+    ax.set_xlabel("Timestep index", fontsize=10)
     ax.set_yticks([0, 1])
     ax.set_yticklabels(["False", "True"])
     ax.grid(True, axis="x", alpha=0.2)
@@ -519,8 +519,9 @@ def _make_trajectory_status_plot(traj, timestep, min_seq_len):
     title = "no_teacher_action | step {} / {}".format(timestep, total_steps - 1)
     if is_valid_start:
         title += " | VALID START"
-    ax.set_title(title)
-    ax.legend(loc="upper right", fontsize=8)
+    ax.set_title(title, fontsize=11)
+    ax.tick_params(axis="both", labelsize=9)
+    ax.legend(loc="upper right", fontsize=9)
 
     # Add saved timestep annotation if the stored timestep is not the same as index.
     saved_timestep = int(timesteps[timestep]) if len(timesteps) else timestep
@@ -584,29 +585,39 @@ def _compose_video_frame(gallery_items, frame_label, status_plot=None):
 
     Top: selected observation images.
     Bottom: trajectory-status plot with the moving timestep cursor.
+
+    Important: do NOT downscale the status plot to the image width. The plot
+    contains tick labels and a legend, so preserving its native width makes the
+    generated MP4 much more readable.
     """
+    small_text_y = 3
+
     if not gallery_items:
         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))
+        draw.text((8, small_text_y), "No selected image keys", fill=(255, 255, 255))
     else:
         pil_images = []
         for img, label in gallery_items:
             pil_img = Image.fromarray(np.asarray(img, dtype=np.uint8)).convert("RGB")
-            label_h = 24
+
+            # Keep the image-key caption compact; large captions waste video space.
+            label_h = 16
             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))
+            draw.text((4, small_text_y), str(label), fill=(220, 220, 220))
             pil_images.append(panel)
 
         gap = 8
-        top_h = 28
+        top_h = 18
         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))
+
+        # Compact frame label above the image panels.
+        draw.text((6, small_text_y), frame_label, fill=(220, 220, 220))
 
         x = 0
         for im in pil_images:
@@ -616,15 +627,22 @@ def _compose_video_frame(gallery_items, frame_label, status_plot=None):
     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)
+        # Preserve the status plot resolution. If needed, pad the observation
+        # canvas to the same width and center it above the plot.
+        final_w = max(obs_canvas.width, status_img.width)
+        if obs_canvas.width < final_w:
+            padded_obs = Image.new("RGB", (final_w, obs_canvas.height), color=(0, 0, 0))
+            padded_obs.paste(obs_canvas, ((final_w - obs_canvas.width) // 2, 0))
+            obs_canvas = padded_obs
+        elif status_img.width < final_w:
+            padded_status = Image.new("RGB", (final_w, status_img.height), color=(255, 255, 255))
+            padded_status.paste(status_img, ((final_w - status_img.width) // 2, 0))
+            status_img = padded_status
 
         gap_h = 8
         canvas = Image.new(
             "RGB",
-            (max(obs_canvas.width, status_img.width), obs_canvas.height + gap_h + status_img.height),
+            (final_w, obs_canvas.height + gap_h + status_img.height),
             color=(0, 0, 0),
         )
         canvas.paste(obs_canvas, (0, 0))
@@ -632,6 +650,7 @@ def _compose_video_frame(gallery_items, frame_label, status_plot=None):
     else:
         canvas = obs_canvas
 
+    # Many MP4 encoders prefer dimensions divisible by 16.
     pad_w = int(np.ceil(canvas.width / 16.0) * 16)
     pad_h = int(np.ceil(canvas.height / 16.0) * 16)
     if pad_w != canvas.width or pad_h != canvas.height:
@@ -884,9 +903,9 @@ 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")
+        trajectory_video = gr.Video(label="Trajectory video: smooth browser-side playback")
         status_plot = gr.Image(label="Trajectory status: no_teacher_action and valid starts", type="numpy")
         info = gr.Textbox(label="Frame info", lines=16)