Update app.py

app.py

@@ -1,38 +1,37 @@
-"""
-Standalone Hugging Face Space viewer for TrajectoryBuffer-style HDF5 files.
-
-Best-practice version:
-- No dependency on your local TrajectoryBuffer Python class.
-- Dataset preset + custom dataset support.
-- Robust HDF5 schema detection for root-level episode_XXXX groups.
-- Auto-detect image keys from each trajectory's observation group.
-- Avoids fragile multiline f-strings in UI status text.
-- Uses slider.release() for timestep rendering to reduce image flicker.
-
-requirements.txt:
-    gradio
-    huggingface_hub
-    h5py
-    numpy
-    pillow
-    matplotlib
-
-Optional:
-    opencv-python-headless
-"""
 
+# Standalone Hugging Face Space viewer for TrajectoryBuffer-style HDF5 files.
+#
+# requirements.txt:
+# gradio
+# huggingface_hub
+# h5py
+# numpy
+# pillow
+# matplotlib
+# imageio
+# imageio-ffmpeg
+#
+# Optional:
+# opencv-python-headless
+
+import os
 import re
+import tempfile
 from functools import lru_cache
 
 import gradio as gr
 import h5py
 import matplotlib
-
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 import numpy as np
 from huggingface_hub import hf_hub_download
-from PIL import Image
+from PIL import Image, ImageDraw
+
+try:
+    import imageio.v2 as imageio
+except Exception:
+    imageio = None
 
 try:
     import cv2
@@ -40,9 +39,6 @@ except Exception:
     cv2 = None
 
 
-# -----------------------------------------------------------------------------
-# Dataset presets
-# -----------------------------------------------------------------------------
 DATASET_PRESETS = {
     "Robosuite Square 20260409": {
         "repo_id": "Zhaoting123/Robosuite_Square_image_abs_with_state",
@@ -67,7 +63,6 @@ DATASET_PRESETS = {
 DEFAULT_PRESET = "Robosuite Square 20260409"
 REPO_TYPE = "dataset"
 DEFAULT_CHUNK_LEN = 16
-
 PREFERRED_IMAGE_KEYS = [
     "image1",
     "image2",
@@ -76,16 +71,11 @@ PREFERRED_IMAGE_KEYS = [
     "front_image",
     "wrist_image",
 ]
-
 IMAGE_KEY_HINTS = ["rgb", "image", "img", "camera", "cam"]
 
 
-# -----------------------------------------------------------------------------
-# Dataset resolution and cache helpers
-# -----------------------------------------------------------------------------
 def resolve_dataset(preset_name, custom_repo_id=None, custom_filename=None):
     preset_name = preset_name or DEFAULT_PRESET
-
     if preset_name == "Custom":
         repo_id = str(custom_repo_id or "").strip()
         filename = str(custom_filename or "").strip()
@@ -93,20 +83,13 @@ def resolve_dataset(preset_name, custom_repo_id=None, custom_filename=None):
             raise ValueError("For Custom mode, provide both repo_id and HDF5 filename/path.")
         return repo_id, filename
 
-    if preset_name not in DATASET_PRESETS:
-        preset_name = DEFAULT_PRESET
-
-    item = DATASET_PRESETS[preset_name]
+    item = DATASET_PRESETS.get(preset_name, DATASET_PRESETS[DEFAULT_PRESET])
     return item["repo_id"], item["filename"]
 
 
 @lru_cache(maxsize=8)
 def get_local_hdf5_path(repo_id, filename):
-    return hf_hub_download(
-        repo_id=repo_id,
-        filename=filename,
-        repo_type=REPO_TYPE,
-    )
+    return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=REPO_TYPE)
 
 
 def _natural_sort_key(name):
@@ -118,39 +101,21 @@ def _natural_sort_key(name):
 
 @lru_cache(maxsize=8)
 def get_trajectory_keys(repo_id, filename):
-    """Return ordered trajectory group paths."""
     path = get_local_hdf5_path(repo_id, filename)
-
     with h5py.File(path, "r") as f:
-        # Your TrajectoryBuffer format:
-        #   /episode_0000
-        #   /episode_0001
         root_episode_keys = [
-            key
-            for key in f.keys()
+            key for key in f.keys()
             if isinstance(f[key], h5py.Group) and str(key).startswith("episode_")
         ]
         if root_episode_keys:
             return tuple(sorted(root_episode_keys, key=_natural_sort_key))
 
-        # Robomimic-style fallback:
-        #   /data/demo_0
-        #   /data/demo_1
         if "data" in f and isinstance(f["data"], h5py.Group):
             data_group = f["data"]
-            keys = [
-                key
-                for key in data_group.keys()
-                if isinstance(data_group[key], h5py.Group)
-            ]
+            keys = [key for key in data_group.keys() if isinstance(data_group[key], h5py.Group)]
             return tuple("data/" + key for key in sorted(keys, key=_natural_sort_key))
 
-        # Generic root-level group fallback.
-        keys = [
-            key
-            for key in f.keys()
-            if isinstance(f[key], h5py.Group)
-        ]
+        keys = [key for key in f.keys() if isinstance(f[key], h5py.Group)]
         return tuple(sorted(keys, key=_natural_sort_key))
 
 
@@ -169,9 +134,7 @@ def inspect_hdf5_tree(preset_name, custom_repo_id, custom_filename, max_lines=18
             if len(lines) >= max_lines:
                 return
             if isinstance(obj, h5py.Dataset):
-                lines.append(
-                    "DATASET {} shape={} dtype={}".format(name, obj.shape, obj.dtype)
-                )
+                lines.append("DATASET {} shape={} dtype={}".format(name, obj.shape, obj.dtype))
             elif isinstance(obj, h5py.Group):
                 lines.append("GROUP   {}".format(name))
 
@@ -179,15 +142,9 @@ def inspect_hdf5_tree(preset_name, custom_repo_id, custom_filename, max_lines=18
 
     if len(lines) >= max_lines:
         lines.append("...")
-
-    if not lines:
-        return "No HDF5 contents found."
-    return chr(10).join(lines)
+    return "\n".join(lines) if lines else "No HDF5 contents found."
 
 
-# -----------------------------------------------------------------------------
-# HDF5 loading helpers
-# -----------------------------------------------------------------------------
 def _read_dataset_value(dataset):
     value = dataset[()]
     if isinstance(value, bytes):
@@ -213,7 +170,6 @@ def _find_first_key(mapping, candidate_keys):
 
 
 def _infer_time_length(data):
-    """Infer trajectory length from common TrajectoryBuffer fields."""
     for key in ["timesteps", "dones", "robot_actions", "teacher_actions", "actions"]:
         if key in data:
             arr = np.asarray(data[key])
@@ -235,7 +191,6 @@ def _infer_time_length(data):
         if lengths:
             values, counts = np.unique(lengths, return_counts=True)
             return int(values[np.argmax(counts)])
-
     return 1
 
 
@@ -248,7 +203,6 @@ def _slice_time(value, t, T):
 
 @lru_cache(maxsize=64)
 def load_traj(repo_id, filename, traj_id):
-    """Load one trajectory as list[dict]."""
     traj_keys = get_trajectory_keys(repo_id, filename)
     if not traj_keys:
         return []
@@ -258,8 +212,7 @@ def load_traj(repo_id, filename, traj_id):
     path = get_local_hdf5_path(repo_id, filename)
 
     with h5py.File(path, "r") as f:
-        group = f[traj_key]
-        data = _read_group_recursive(group)
+        data = _read_group_recursive(f[traj_key])
 
     T = _infer_time_length(data)
 
@@ -281,66 +234,41 @@ def load_traj(repo_id, filename, traj_id):
 
     traj = []
     for t in range(T):
-        obs_t = {}
-        for key, value in obs_all.items():
-            obs_t[key] = _slice_time(value, t, T)
+        obs_t = {key: _slice_time(value, t, T) for key, value in obs_all.items()}
 
         default_action = np.zeros(1, dtype=np.float32)
         if action_key is not None:
             default_action = _slice_time(data[action_key], t, T)
 
-        teacher_action = default_action
-        if teacher_key is not None:
-            teacher_action = _slice_time(data[teacher_key], t, T)
-
-        robot_action = default_action
-        if robot_key is not None:
-            robot_action = _slice_time(data[robot_key], t, T)
-
-        no_teacher = False
-        if no_teacher_key is not None:
-            no_teacher = _slice_time(data[no_teacher_key], t, T)
-
-        no_robot = False
-        if no_robot_key is not None:
-            no_robot = _slice_time(data[no_robot_key], t, T)
-
-        done = False
-        if done_key is not None:
-            done = _slice_time(data[done_key], t, T)
+        teacher_action = _slice_time(data[teacher_key], t, T) if teacher_key else default_action
+        robot_action = _slice_time(data[robot_key], t, T) if robot_key else default_action
+        no_teacher = _slice_time(data[no_teacher_key], t, T) if no_teacher_key else False
+        no_robot = _slice_time(data[no_robot_key], t, T) if no_robot_key else False
+        done = _slice_time(data[done_key], t, T) if done_key else False
+        if_success = _slice_time(data[success_key], t, T) if success_key else False
 
         timestep = t
         if timestep_key is not None:
             timestep_arr = _slice_time(data[timestep_key], t, T)
             timestep = int(np.asarray(timestep_arr).reshape(-1)[0])
 
-        if_success = False
-        if success_key is not None:
-            if_success = _slice_time(data[success_key], t, T)
-
-        traj.append(
-            {
-                "obs": obs_t,
-                "robot_action": np.asarray(robot_action),
-                "teacher_action": np.asarray(teacher_action),
-                "done": bool(np.asarray(done).reshape(-1)[0]),
-                "timestep": timestep,
-                "no_robot_action": bool(np.asarray(no_robot).reshape(-1)[0]),
-                "no_teacher_action": bool(np.asarray(no_teacher).reshape(-1)[0]),
-                "episode_id": traj_key,
-                "if_success": bool(np.asarray(if_success).reshape(-1)[0]),
-            }
-        )
+        traj.append({
+            "obs": obs_t,
+            "robot_action": np.asarray(robot_action),
+            "teacher_action": np.asarray(teacher_action),
+            "done": bool(np.asarray(done).reshape(-1)[0]),
+            "timestep": timestep,
+            "no_robot_action": bool(np.asarray(no_robot).reshape(-1)[0]),
+            "no_teacher_action": bool(np.asarray(no_teacher).reshape(-1)[0]),
+            "episode_id": traj_key,
+            "if_success": bool(np.asarray(if_success).reshape(-1)[0]),
+        })
 
     return traj
 
 
-# -----------------------------------------------------------------------------
-# Image and plotting helpers
-# -----------------------------------------------------------------------------
 def _extract_latest_obs_value(value):
     arr = np.asarray(value)
-    # Per-timestep stacked observation commonly has shape [obs_T, C, H, W].
     if arr.ndim >= 1 and arr.shape[0] in (1, 2, 3, 4):
         return arr[-1]
     return arr
@@ -365,8 +293,6 @@ def _float_img_to_uint8(img):
     arr_min = float(np.nanmin(arr))
     arr_max = float(np.nanmax(arr))
 
-    # TrajectoryBuffer saves float images originally in [-1, 1] as uint8.
-    # But for compatibility, handle float [-1, 1], [0, 1], and [0, 255].
     if arr_min >= -1.01 and arr_max <= 1.01:
         if arr_min < 0.0:
             arr = (arr + 1.0) * 0.5
@@ -395,16 +321,10 @@ def _extract_display_image(value, reverse_channels=False):
     if img.ndim != 3:
         raise ValueError("Unsupported image shape: {}".format(img.shape))
 
-    if img.dtype == np.uint8:
-        out = img.copy()
-    else:
-        out = _float_img_to_uint8(img)
+    out = img.copy() if img.dtype == np.uint8 else _float_img_to_uint8(img)
 
-    # Browser display expects RGB. Your current data appears RGB already,
-    # so default reverse_channels=False.
     if reverse_channels and out.shape[-1] == 3:
         out = out[..., ::-1]
-
     return out
 
 
@@ -437,7 +357,6 @@ def _extract_mixed_action_chunk(traj, start_idx, chunk_length):
 
     if not chunk:
         return None, ""
-
     return np.stack(chunk, axis=0), "".join(sources)
 
 
@@ -451,7 +370,6 @@ def _extract_robot_action_chunk(traj, start_idx, chunk_length):
 
     if not chunk:
         return None
-
     return np.stack(chunk, axis=0)
 
 
@@ -505,25 +423,8 @@ 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.
-    """
+def get_cached_gallery_items(repo_id, filename, traj_id, timestep, image_keys_tuple, display_scale, reverse_channels):
     traj = load_traj(repo_id, filename, int(traj_id))
     timestep = int(np.clip(int(timestep), 0, len(traj) - 1))
     obs = traj[timestep].get("obs", {})
@@ -535,10 +436,7 @@ def get_cached_gallery_items(
             warnings.append("Missing image key: {}".format(key))
             continue
         try:
-            img = _extract_display_image(
-                obs[key],
-                reverse_channels=bool(reverse_channels),
-            )
+            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:
@@ -556,17 +454,7 @@ def get_cached_action_plot(repo_id, filename, traj_id, timestep, 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."""
+def preload_current_trajectory(preset_name, custom_repo_id, custom_filename, traj_id, image_keys, chunk_len, display_scale, reverse_channels):
     repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
     n_traj = get_num_trajectories(repo_id, filename)
     if n_traj == 0:
@@ -585,27 +473,106 @@ def preload_current_trajectory(
 
     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_gallery_items(repo_id, filename, traj_id, t, image_keys_tuple, float(display_scale), bool(reverse_channels))
         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")
+    status += "\nFrames cached: {}".format(total)
+    status += "\nImage keys: {}".format(", ".join(image_keys_tuple) if image_keys_tuple else "none")
     return status
 
 
-# -----------------------------------------------------------------------------
-# Gradio callbacks
-# -----------------------------------------------------------------------------
+def _compose_video_frame(gallery_items, frame_label):
+    if not gallery_items:
+        canvas = Image.new("RGB", (640, 360), color=(20, 20, 20))
+        draw = ImageDraw.Draw(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
+
+    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:
+        padded = Image.new("RGB", (pad_w, pad_h), color=(0, 0, 0))
+        padded.paste(canvas, (0, 0))
+        canvas = padded
+
+    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):
+    if imageio is None:
+        return None, "Video export requires imageio and imageio-ffmpeg in requirements.txt."
+
+    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 None, "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 None, "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)
+
+    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)),
+    )
+
+    writer = imageio.get_writer(out_path, fps=float(fps), codec="libx264", quality=8)
+    try:
+        for t in range(len(traj)):
+            gallery_items, _warnings = get_cached_gallery_items(
+                repo_id,
+                filename,
+                traj_id,
+                t,
+                image_keys_tuple,
+                float(display_scale),
+                bool(reverse_channels),
+            )
+            label = "trajectory {} | frame {}/{}".format(traj_id, t, len(traj) - 1)
+            frame = _compose_video_frame(gallery_items, label)
+            writer.append_data(frame)
+    finally:
+        writer.close()
+
+    status = "Built trajectory video"
+    status += "\nTrajectory: {}".format(traj_id)
+    status += "\nFrames: {} | FPS: {}".format(len(traj), fps)
+    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:
@@ -641,7 +608,7 @@ def update_after_dataset_change(preset_name, custom_repo_id, custom_filename):
 
     if n_traj == 0:
         status = "Loaded `{}` / `{}`".format(repo_id, filename)
-        status = status + chr(10) + "Detected trajectories: 0"
+        status += "\nDetected trajectories: 0"
         return (
             gr.update(maximum=1, value=0),
             gr.update(maximum=1, value=0),
@@ -653,7 +620,7 @@ def update_after_dataset_change(preset_name, custom_repo_id, custom_filename):
     traj = load_traj(repo_id, filename, 0)
 
     status = "Loaded `{}` / `{}`".format(repo_id, filename)
-    status = status + chr(10) + "Detected trajectories: {}".format(n_traj)
+    status += "\nDetected trajectories: {}".format(n_traj)
 
     return (
         gr.update(maximum=max(n_traj - 1, 1), value=0),
@@ -679,17 +646,7 @@ def update_after_traj_change(preset_name, custom_repo_id, custom_filename, traj_
     )
 
 
-def render_frame(
-    preset_name,
-    custom_repo_id,
-    custom_filename,
-    traj_id,
-    timestep,
-    image_keys,
-    chunk_len,
-    display_scale,
-    reverse_channels,
-):
+def render_frame(preset_name, custom_repo_id, custom_filename, traj_id, timestep, image_keys, chunk_len, display_scale, reverse_channels):
     repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
     n_traj = get_num_trajectories(repo_id, filename)
 
@@ -711,27 +668,14 @@ def render_frame(
         image_keys = [image_keys]
 
     step = traj[timestep]
-    obs = step.get("obs", {})
-
     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),
+        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,
-    )
+    action_plot, source_mask = get_cached_action_plot(repo_id, filename, traj_id, timestep, chunk_len)
 
     info_lines = [
         "dataset: {} / {}".format(repo_id, filename),
@@ -756,12 +700,9 @@ def render_frame(
         info_lines.append("Image warnings:")
         info_lines.extend(warnings)
 
-    return gallery_items, action_plot, chr(10).join(info_lines)
+    return gallery_items, action_plot, "\n".join(info_lines)
 
 
-# -----------------------------------------------------------------------------
-# App
-# -----------------------------------------------------------------------------
 def build_app():
     repo_id, filename = resolve_dataset(DEFAULT_PRESET)
 
@@ -774,8 +715,7 @@ def build_app():
         first_keys = []
         startup_warning = repr(exc)
 
-    default_status = "Loaded default dataset" + chr(10)
-    default_status += "Detected trajectories: {}".format(n_traj)
+    default_status = "Loaded default dataset\nDetected trajectories: {}".format(n_traj)
 
     with gr.Blocks(title="HDF5 Trajectory Viewer") as demo:
         gr.Markdown(
@@ -792,82 +732,31 @@ def build_app():
                 value=DEFAULT_PRESET,
                 label="Dataset preset",
             )
-            custom_repo_id = gr.Textbox(
-                value="",
-                label="Custom repo_id, e.g. Zhaoting123/InsertT",
-                visible=False,
-            )
-            custom_filename = gr.Textbox(
-                value="",
-                label="Custom HDF5 path in repo",
-                visible=False,
-            )
+            custom_repo_id = gr.Textbox(value="", label="Custom repo_id, e.g. Zhaoting123/InsertT", visible=False)
+            custom_filename = gr.Textbox(value="", label="Custom HDF5 path in repo", visible=False)
 
-        dataset_status = gr.Textbox(
-            label="Dataset status",
-            lines=2,
-            value=default_status,
-            interactive=False,
-        )
+        dataset_status = gr.Textbox(label="Dataset status", lines=2, value=default_status, interactive=False)
 
         with gr.Row():
-            traj_slider = gr.Slider(
-                minimum=0,
-                maximum=max(n_traj - 1, 1),
-                value=0,
-                step=1,
-                label="Trajectory index",
-            )
-            timestep_slider = gr.Slider(
-                minimum=0,
-                maximum=1,
-                value=0,
-                step=1,
-                label="Timestep",
-            )
+            traj_slider = gr.Slider(minimum=0, maximum=max(n_traj - 1, 1), value=0, step=1, label="Trajectory index")
+            timestep_slider = gr.Slider(minimum=0, maximum=1, value=0, step=1, label="Timestep")
 
         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="Action chunk length",
-            )
-            display_scale = gr.Slider(
-                minimum=1,
-                maximum=10,
-                value=4,
-                step=1,
-                label="Image display scale",
-            )
-            reverse_channels = gr.Checkbox(
-                value=False,
-                label="Reverse channels BGR↔RGB",
-            )
+            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="Action chunk length")
+            display_scale = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Image display scale")
+            reverse_channels = gr.Checkbox(value=False, label="Reverse channels BGR↔RGB")
 
         with gr.Row():
             render_btn = gr.Button("Render frame", variant="primary")
             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")
 
-        preload_status = gr.Textbox(
-            label="Preload status",
-            lines=3,
-            value="Not preloaded yet.",
-            interactive=False,
-        )
+        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",
-        )
+        gallery = gr.Gallery(label="Observation images", columns=2, height="auto", object_fit="contain")
         action_plot = gr.Image(label="Action chunk plot", type="numpy")
         info = gr.Textbox(label="Frame info", lines=16)
 
@@ -875,7 +764,6 @@ def build_app():
             inspect_btn = gr.Button("Inspect HDF5 structure")
             hdf5_tree = gr.Textbox(lines=24, label="HDF5 tree")
 
-        # Dataset selection and custom-field visibility.
         preset.change(
             fn=update_custom_visibility,
             inputs=preset,
@@ -886,17 +774,7 @@ def build_app():
             outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
         ).then(
             fn=render_frame,
-            inputs=[
-                preset,
-                custom_repo_id,
-                custom_filename,
-                traj_slider,
-                timestep_slider,
-                image_keys,
-                chunk_len,
-                display_scale,
-                reverse_channels,
-            ],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
             outputs=[gallery, action_plot, info],
         )
 
@@ -917,85 +795,41 @@ def build_app():
             outputs=[timestep_slider, image_keys],
         ).then(
             fn=render_frame,
-            inputs=[
-                preset,
-                custom_repo_id,
-                custom_filename,
-                traj_slider,
-                timestep_slider,
-                image_keys,
-                chunk_len,
-                display_scale,
-                reverse_channels,
-            ],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
             outputs=[gallery, action_plot, info],
         )
 
-        # Render only after releasing the timestep slider, reducing flicker.
         timestep_slider.release(
             fn=render_frame,
-            inputs=[
-                preset,
-                custom_repo_id,
-                custom_filename,
-                traj_slider,
-                timestep_slider,
-                image_keys,
-                chunk_len,
-                display_scale,
-                reverse_channels,
-            ],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
             outputs=[gallery, action_plot, info],
         )
 
         for widget in [image_keys, chunk_len, display_scale, 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,
-                ],
+                inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
                 outputs=[gallery, action_plot, info],
             )
 
         render_btn.click(
             fn=render_frame,
-            inputs=[
-                preset,
-                custom_repo_id,
-                custom_filename,
-                traj_slider,
-                timestep_slider,
-                image_keys,
-                chunk_len,
-                display_scale,
-                reverse_channels,
-            ],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
             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,
-            ],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, image_keys, chunk_len, display_scale, reverse_channels],
             outputs=preload_status,
         )
 
+        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],
+            outputs=[trajectory_video, preload_status],
+        )
+
         inspect_btn.click(
             fn=inspect_hdf5_tree,
             inputs=[preset, custom_repo_id, custom_filename],
@@ -1008,17 +842,7 @@ def build_app():
             outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
         ).then(
             fn=render_frame,
-            inputs=[
-                preset,
-                custom_repo_id,
-                custom_filename,
-                traj_slider,
-                timestep_slider,
-                image_keys,
-                chunk_len,
-                display_scale,
-                reverse_channels,
-            ],
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
             outputs=[gallery, action_plot, info],
         )
 
@@ -1026,4 +850,4 @@ def build_app():
 
 
 if __name__ == "__main__":
-    build_app().launch()
+    build_app().launch()