Update app.py

app.py

@@ -68,6 +68,7 @@ DATASET_PRESETS = {
 DEFAULT_PRESET = "Robosuite Square Correction"
 REPO_TYPE = "dataset"
 DEFAULT_CHUNK_LEN = 16
+DEFAULT_DISPLAY_SCALE = 1
 PREFERRED_IMAGE_KEYS = [
     "image1",
     "image2",
@@ -797,7 +798,58 @@ def get_cached_video_status_frame(repo_id, filename, traj_id, timestep, valid_wi
 
     return np.asarray(img)
 
-def build_current_trajectory_video(preset_name, custom_repo_id, custom_filename, traj_id, image_keys, display_scale, reverse_channels, fps, valid_window_len):
+def _draw_status_cursor_on_base(base, bounds, total_steps, timestep):
+    """Fast video status frame: copy one static Matplotlib image and draw cursor.
+
+    This avoids calling the lru-cached per-timestep status frame function during
+    video export. For long trajectories, caching thousands of status images can
+    consume a lot of memory and still requires PIL conversion for every frame.
+    """
+    if base is None:
+        return None
+
+    total_steps = int(max(total_steps, 1))
+    timestep = int(np.clip(int(timestep), 0, total_steps - 1))
+    x0, y0, x1, y1 = [int(v) for v in bounds]
+    denom = max(total_steps - 1, 1)
+    x = int(round(x0 + (x1 - x0) * float(timestep) / float(denom)))
+
+    img = np.asarray(base, dtype=np.uint8).copy()
+
+    # Draw the vertical cursor directly with NumPy. This is much cheaper than
+    # creating a Matplotlib plot for every frame.
+    x_left = max(0, x - 2)
+    x_right = min(img.shape[1], x + 2)
+    y_top = max(0, y0)
+    y_bottom = min(img.shape[0], y1)
+    img[y_top:y_bottom, x_left:x_right, :] = 0
+
+    # Small text label. PIL is used only for the label, not for the whole plot.
+    pil_img = Image.fromarray(img).convert("RGB")
+    draw = ImageDraw.Draw(pil_img)
+    label = "step {}/{}".format(timestep, total_steps - 1)
+    draw.rectangle((x0 + 4, y0 + 4, x0 + 126, y0 + 24), fill=(255, 255, 255))
+    draw.text((x0 + 8, y0 + 7), label, fill=(0, 0, 0))
+    return np.asarray(pil_img)
+
+
+def _get_fast_video_writer(out_path, fps):
+    """Use ffmpeg's ultrafast x264 preset for interactive Spaces exports."""
+    return imageio.get_writer(
+        out_path,
+        fps=float(fps),
+        codec="libx264",
+        macro_block_size=16,
+        ffmpeg_params=[
+            "-preset", "ultrafast",
+            "-crf", "28",
+            "-pix_fmt", "yuv420p",
+            "-movflags", "+faststart",
+        ],
+    )
+
+
+def build_current_trajectory_video(preset_name, custom_repo_id, custom_filename, traj_id, image_keys, display_scale, reverse_channels, fps, valid_window_len, video_stride=1):
     if imageio is None:
         return None, "Video export requires imageio and imageio-ffmpeg in requirements.txt."
 
@@ -817,16 +869,34 @@ def build_current_trajectory_video(preset_name, custom_repo_id, custom_filename,
         image_keys = [image_keys]
     image_keys_tuple = tuple(image_keys)
 
+    video_stride = int(max(1, int(video_stride)))
+    frame_indices = list(range(0, len(traj), video_stride))
+    if frame_indices and frame_indices[-1] != len(traj) - 1:
+        frame_indices.append(len(traj) - 1)
+
     safe_repo = re.sub(r"[^A-Za-z0-9_.-]+", "_", repo_id)
     safe_file = re.sub(r"[^A-Za-z0-9_.-]+", "_", filename)[-80:]
     out_path = os.path.join(
         tempfile.gettempdir(),
-        "trajectory_{}_{}_traj{:04d}_fps{}.mp4".format(safe_repo, safe_file, traj_id, int(fps)),
+        "trajectory_{}_{}_traj{:04d}_fps{}_stride{}.mp4".format(
+            safe_repo, safe_file, traj_id, int(fps), video_stride
+        ),
+    )
+
+    # Build the static status plot once. During export, only draw the cursor.
+    status_base, status_bounds, total_steps = get_video_status_plot_base(
+        repo_id,
+        filename,
+        traj_id,
+        int(valid_window_len),
     )
 
-    writer = imageio.get_writer(out_path, fps=float(fps), codec="libx264", quality=8)
+    writer = _get_fast_video_writer(out_path, fps)
+    written = 0
     try:
-        for t in range(len(traj)):
+        for t in frame_indices:
+            # Use the existing cached image extraction for correctness, but avoid
+            # cached per-timestep status images to reduce memory pressure.
             gallery_items, _warnings = get_cached_gallery_items(
                 repo_id,
                 filename,
@@ -837,26 +907,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)
-            status_plot = get_cached_video_status_frame(
-                repo_id,
-                filename,
-                traj_id,
-                t,
-                int(valid_window_len),
-            )
+            status_plot = _draw_status_cursor_on_base(status_base, status_bounds, total_steps, t)
             frame = _compose_video_frame(gallery_items, label, status_plot=status_plot)
             writer.append_data(frame)
+            written += 1
     finally:
         writer.close()
 
-    status = "Built trajectory video with fast moving status plot"
+    approx_seconds = float(written) / float(max(float(fps), 1.0))
+    status = "Built trajectory video with optimized encoder and status rendering"
     status += "\nTrajectory: {}".format(traj_id)
-    status += "\nFrames: {} | FPS: {}".format(len(traj), fps)
+    status += "\nOriginal timesteps: {} | Written frames: {} | Stride: {}".format(len(traj), written, video_stride)
+    status += "\nFPS: {} | Approx video duration: {:.1f}s".format(fps, approx_seconds)
     status += "\nValid-window length: {}".format(int(valid_window_len))
-    status += "\nSpeedup: static plot rendered once; cursor drawn per frame"
+    status += "\nSpeedups: x264 ultrafast preset; static status plot rendered once; cursor drawn with NumPy/PIL"
     return out_path, status
 
-
 def get_available_image_keys(repo_id, filename, traj_id):
     n_traj = get_num_trajectories(repo_id, filename)
     if n_traj == 0:
@@ -1047,7 +1113,7 @@ def build_app():
         with gr.Row():
             image_keys = gr.CheckboxGroup(choices=first_keys, value=first_keys[:2], label="Image keys")
             chunk_len = gr.Slider(minimum=1, maximum=64, value=DEFAULT_CHUNK_LEN, step=1, label="Valid-window length")
-            display_scale = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Image display scale")
+            display_scale = gr.State(value=DEFAULT_DISPLAY_SCALE)
             reverse_channels = gr.Checkbox(value=get_default_reverse_channels(DEFAULT_PRESET), label="Reverse channels BGR↔RGB")
 
         with gr.Row():
@@ -1055,6 +1121,7 @@ def build_app():
             preload_btn = gr.Button("Preload current trajectory")
             video_btn = gr.Button("Build trajectory video")
             video_fps = gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Video FPS")
+            video_stride = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="Video frame stride")
 
         preload_status = gr.Textbox(label="Preload / video status", lines=4, value="Not preloaded yet.", interactive=False)
 
@@ -1121,7 +1188,7 @@ def build_app():
             outputs=[gallery, status_plot, info],
         )
 
-        for widget in [image_keys, chunk_len, display_scale, reverse_channels]:
+        for widget in [image_keys, chunk_len, reverse_channels]:
             widget.change(
                 fn=render_frame,
                 inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
@@ -1142,7 +1209,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, chunk_len],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, image_keys, display_scale, reverse_channels, video_fps, chunk_len, video_stride],
             outputs=[trajectory_video, preload_status],
         )