Update app.py

app.py

@@ -505,6 +505,104 @@ def _make_action_chunk_plot(mixed_chunk, robot_chunk):
     return image
 
 
+# -----------------------------------------------------------------------------
+# Frame-level render cache
+# -----------------------------------------------------------------------------
+@lru_cache(maxsize=8192)
+def get_cached_gallery_items(
+    repo_id,
+    filename,
+    traj_id,
+    timestep,
+    image_keys_tuple,
+    display_scale,
+    reverse_channels,
+):
+    """Cache decoded/resized observation images for one frame.
+
+    Gradio may briefly clear the image component while a callback is running.
+    This cache makes the callback fast after preloading, which largely removes
+    the black-frame effect when scrubbing.
+    """
+    traj = load_traj(repo_id, filename, int(traj_id))
+    timestep = int(np.clip(int(timestep), 0, len(traj) - 1))
+    obs = traj[timestep].get("obs", {})
+
+    gallery_items = []
+    warnings = []
+    for key in image_keys_tuple:
+        if key not in obs:
+            warnings.append("Missing image key: {}".format(key))
+            continue
+        try:
+            img = _extract_display_image(
+                obs[key],
+                reverse_channels=bool(reverse_channels),
+            )
+            img = _resize_image_for_display(img, float(display_scale))
+            gallery_items.append((img, key))
+        except Exception as exc:
+            warnings.append("{}: {}".format(key, exc))
+
+    return gallery_items, tuple(warnings)
+
+
+@lru_cache(maxsize=8192)
+def get_cached_action_plot(repo_id, filename, traj_id, timestep, chunk_len):
+    traj = load_traj(repo_id, filename, int(traj_id))
+    timestep = int(np.clip(int(timestep), 0, len(traj) - 1))
+    mixed_chunk, source_mask = _extract_mixed_action_chunk(traj, timestep, int(chunk_len))
+    robot_chunk = _extract_robot_action_chunk(traj, timestep, int(chunk_len))
+    return _make_action_chunk_plot(mixed_chunk, robot_chunk), source_mask
+
+
+def preload_current_trajectory(
+    preset_name,
+    custom_repo_id,
+    custom_filename,
+    traj_id,
+    image_keys,
+    chunk_len,
+    display_scale,
+    reverse_channels,
+):
+    """Pre-render all selected observation frames for the current trajectory."""
+    repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
+    n_traj = get_num_trajectories(repo_id, filename)
+    if n_traj == 0:
+        return "No trajectories found."
+
+    traj_id = int(np.clip(int(traj_id), 0, n_traj - 1))
+    traj = load_traj(repo_id, filename, traj_id)
+    if not traj:
+        return "Trajectory could not be loaded."
+
+    if image_keys is None:
+        image_keys = []
+    if isinstance(image_keys, str):
+        image_keys = [image_keys]
+    image_keys_tuple = tuple(image_keys)
+
+    total = len(traj)
+    for t in range(total):
+        get_cached_gallery_items(
+            repo_id,
+            filename,
+            traj_id,
+            t,
+            image_keys_tuple,
+            float(display_scale),
+            bool(reverse_channels),
+        )
+        # The action plot is usually the slowest part. Preload it too.
+        get_cached_action_plot(repo_id, filename, traj_id, t, int(chunk_len))
+
+    status = "Preloaded trajectory {}".format(traj_id)
+    status += chr(10) + "Frames cached: {}".format(total)
+    status += chr(10) + "Image keys: {}".format(", ".join(image_keys_tuple) if image_keys_tuple else "none")
+    return status
+
+
 # -----------------------------------------------------------------------------
 # Gradio callbacks
 # -----------------------------------------------------------------------------
@@ -615,25 +713,25 @@ def render_frame(
     step = traj[timestep]
     obs = step.get("obs", {})
 
-    gallery_items = []
-    warnings = []
-    for key in image_keys:
-        if key not in obs:
-            warnings.append("Missing image key: {}".format(key))
-            continue
-        try:
-            img = _extract_display_image(
-                obs[key],
-                reverse_channels=bool(reverse_channels),
-            )
-            img = _resize_image_for_display(img, display_scale)
-            gallery_items.append((img, key))
-        except Exception as exc:
-            warnings.append("{}: {}".format(key, exc))
-
-    mixed_chunk, source_mask = _extract_mixed_action_chunk(traj, timestep, chunk_len)
-    robot_chunk = _extract_robot_action_chunk(traj, timestep, chunk_len)
-    action_plot = _make_action_chunk_plot(mixed_chunk, robot_chunk)
+    image_keys_tuple = tuple(image_keys)
+    gallery_items, warnings_tuple = get_cached_gallery_items(
+        repo_id,
+        filename,
+        traj_id,
+        timestep,
+        image_keys_tuple,
+        display_scale,
+        bool(reverse_channels),
+    )
+    warnings = list(warnings_tuple)
+
+    action_plot, source_mask = get_cached_action_plot(
+        repo_id,
+        filename,
+        traj_id,
+        timestep,
+        chunk_len,
+    )
 
     info_lines = [
         "dataset: {} / {}".format(repo_id, filename),
@@ -753,7 +851,16 @@ def build_app():
                 label="Reverse channels BGR↔RGB",
             )
 
-        render_btn = gr.Button("Render frame", variant="primary")
+        with gr.Row():
+            render_btn = gr.Button("Render frame", variant="primary")
+            preload_btn = gr.Button("Preload current trajectory")
+
+        preload_status = gr.Textbox(
+            label="Preload status",
+            lines=3,
+            value="Not preloaded yet.",
+            interactive=False,
+        )
 
         gallery = gr.Gallery(
             label="Observation images",
@@ -874,6 +981,21 @@ def build_app():
             outputs=[gallery, action_plot, info],
         )
 
+        preload_btn.click(
+            fn=preload_current_trajectory,
+            inputs=[
+                preset,
+                custom_repo_id,
+                custom_filename,
+                traj_slider,
+                image_keys,
+                chunk_len,
+                display_scale,
+                reverse_channels,
+            ],
+            outputs=preload_status,
+        )
+
         inspect_btn.click(
             fn=inspect_hdf5_tree,
             inputs=[preset, custom_repo_id, custom_filename],