Update app.py

app.py

@@ -40,15 +40,32 @@ except Exception:
 
 
 # -----------------------------------------------------------------------------
-# EDIT THESE FOR YOUR DATASET
+# Dataset presets.
+# The same Space can visualize multiple HDF5 files by changing repo_id + filename.
 # -----------------------------------------------------------------------------
-REPO_ID = "Zhaoting123/Robosuite_Square_image_abs_with_state"
+DATASET_PRESETS = {
+    "Robosuite Square 20260409": {
+        "repo_id": "Zhaoting123/Robosuite_Square_image_abs_with_state",
+        "filename": (
+            "20260409_205051_Diffusion_CLIC_intervention_Circular_square_image_abs_"
+            "Ta16_offlineFalse_Scale0.01/trajectory_buffer_0.hdf5"
+        ),
+    },
+    "Robosuite Square 20260410": {
+        "repo_id": "Zhaoting123/Robosuite_Square_image_abs_with_state",
+        "filename": (
+            "20260410_205606_Diffusion_CLIC_intervention_Circular_square_image_abs_"
+            "Ta16_offlineFalse_Scale0.01/trajectory_buffer_0.hdf5"
+        ),
+    },
+    "InsertT Nov10 noisy": {
+        "repo_id": "Zhaoting123/InsertT",
+        "filename": "trajectory_buffer_Nov10_demo_noisy.hdf5",
+    },
+}
+
+DEFAULT_PRESET = "Robosuite Square 20260409"
 REPO_TYPE = "dataset"
-HDF5_FILENAME = (
-    "20260409_205051_Diffusion_CLIC_intervention_Circular_square_image_abs_"
-    "Ta16_offlineFalse_Scale0.01/trajectory_buffer_0.hdf5"
-)
-
 DEFAULT_CHUNK_LEN = 16
 PREFERRED_IMAGE_KEYS = [
     "image1",
@@ -61,11 +78,34 @@ PREFERRED_IMAGE_KEYS = [
 # -----------------------------------------------------------------------------
 # HDF5 helpers
 # -----------------------------------------------------------------------------
-@lru_cache(maxsize=1)
-def get_local_hdf5_path():
+def _clear_dataset_caches():
+    get_local_hdf5_path.cache_clear()
+    get_trajectory_keys.cache_clear()
+    get_num_trajectories.cache_clear()
+    load_traj.cache_clear()
+
+
+def resolve_dataset(preset_name, custom_repo_id=None, custom_filename=None):
+    """Return (repo_id, filename) from a preset or custom fields."""
+    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()
+        if not repo_id or not filename:
+            raise ValueError("For Custom, 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]
+    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=HDF5_FILENAME,
+        repo_id=repo_id,
+        filename=filename,
         repo_type=REPO_TYPE,
     )
 
@@ -75,10 +115,10 @@ def _natural_sort_key(name):
     return (0, int(m.group(1))) if m else (1, str(name))
 
 
-@lru_cache(maxsize=1)
-def get_trajectory_keys():
+@lru_cache(maxsize=8)
+def get_trajectory_keys(repo_id, filename):
     """Detect trajectory groups in common HDF5 layouts."""
-    path = get_local_hdf5_path()
+    path = get_local_hdf5_path(repo_id, filename)
     with h5py.File(path, "r") as f:
         # Your TrajectoryBuffer saves root-level groups:
         #   /episode_0000
@@ -106,14 +146,15 @@ def get_trajectory_keys():
     return tuple(f"{prefix}/{k}" if prefix else k for k in group_keys)
 
 
-@lru_cache(maxsize=1)
-def get_num_trajectories():
-    return max(len(get_trajectory_keys()), 1)
+@lru_cache(maxsize=8)
+def get_num_trajectories(repo_id, filename):
+    return max(len(get_trajectory_keys(repo_id, filename)), 1)
 
 
-def inspect_hdf5_tree(max_lines=160):
+def inspect_hdf5_tree(preset_name, custom_repo_id, custom_filename, max_lines=160):
     """Show the HDF5 tree for debugging inside the Space."""
-    path = get_local_hdf5_path()
+    repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
+    path = get_local_hdf5_path(repo_id, filename)
     lines = []
     with h5py.File(path, "r") as f:
         def visitor(name, obj):
@@ -184,8 +225,8 @@ def _infer_time_length(data):
     return int(values[np.argmax(counts)])
 
 
-@lru_cache(maxsize=32)
-def load_traj(traj_id):
+@lru_cache(maxsize=64)
+def load_traj(repo_id, filename, traj_id):
     """Load one trajectory as a list of step dictionaries.
 
     Output step format:
@@ -198,8 +239,8 @@ def load_traj(traj_id):
         "no_robot_action": bool,
       }
     """
-    path = get_local_hdf5_path()
-    traj_keys = get_trajectory_keys()
+    path = get_local_hdf5_path(repo_id, filename)
+    traj_keys = get_trajectory_keys(repo_id, filename)
     if not traj_keys:
         return []
 
@@ -415,10 +456,10 @@ def _make_action_chunk_plot(mixed_chunk, robot_chunk=None):
     return img
 
 
-def get_available_image_keys(traj_id):
-    n_traj = get_num_trajectories()
+def get_available_image_keys(repo_id, filename, traj_id):
+    n_traj = get_num_trajectories(repo_id, filename)
     traj_id = int(np.clip(int(traj_id), 0, max(n_traj - 1, 0)))
-    traj = load_traj(traj_id)
+    traj = load_traj(repo_id, filename, traj_id)
     if not traj:
         return []
 
@@ -447,11 +488,12 @@ def get_available_image_keys(traj_id):
 # -----------------------------------------------------------------------------
 # Gradio callbacks
 # -----------------------------------------------------------------------------
-def update_after_traj_change(traj_id):
-    n_traj = get_num_trajectories()
+def update_after_traj_change(preset_name, custom_repo_id, custom_filename, traj_id):
+    repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
+    n_traj = get_num_trajectories(repo_id, filename)
     traj_id = int(np.clip(int(traj_id), 0, max(n_traj - 1, 0)))
-    traj = load_traj(traj_id)
-    image_keys = get_available_image_keys(traj_id)
+    traj = load_traj(repo_id, filename, traj_id)
+    image_keys = get_available_image_keys(repo_id, filename, traj_id)
     max_step = max(len(traj) - 1, 0)
     slider_max = max(max_step, 1)  # Gradio requires min < max.
     return (
@@ -460,10 +502,11 @@ def update_after_traj_change(traj_id):
     )
 
 
-def render_frame(traj_id, timestep, image_keys, chunk_len, display_scale, reverse_channels):
-    n_traj = get_num_trajectories()
+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)
     traj_id = int(np.clip(int(traj_id), 0, max(n_traj - 1, 0)))
-    traj = load_traj(traj_id)
+    traj = load_traj(repo_id, filename, traj_id)
 
     if not traj:
         return [], None, "No trajectory could be loaded. Open the HDF5 debug panel to inspect the file layout."
@@ -525,8 +568,9 @@ def render_frame(traj_id, timestep, image_keys, chunk_len, display_scale, revers
 # -----------------------------------------------------------------------------
 def build_app():
     try:
-        n_traj = get_num_trajectories()
-        first_keys = get_available_image_keys(0)
+        repo_id, filename = resolve_dataset(DEFAULT_PRESET)
+        n_traj = get_num_trajectories(repo_id, filename)
+        first_keys = get_available_image_keys(repo_id, filename, 0)
         startup_error = None
     except Exception as exc:
         n_traj = 1
@@ -537,7 +581,7 @@ def build_app():
         gr.Markdown(
             "# HDF5 Trajectory Viewer\n"
             "Standalone viewer: no local `TrajectoryBuffer` dependency.\n\n"
-            f"Detected trajectories: **{n_traj}**"
+            f"Default dataset detected trajectories: **{n_traj}**"
         )
 
         if startup_error is not None:
@@ -546,6 +590,23 @@ def build_app():
                 f"```text\n{startup_error}\n```"
             )
 
+        with gr.Row():
+            preset = gr.Dropdown(
+                choices=list(DATASET_PRESETS.keys()) + ["Custom"],
+                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,
+            )
+
         with gr.Row():
             traj_slider = gr.Slider(
                 minimum=0,
@@ -583,7 +644,7 @@ def build_app():
                 label="Image display scale",
             )
             reverse_channels = gr.Checkbox(
-                value=True,
+                value=False,
                 label="Reverse channels BGR↔RGB",
             )
 
@@ -601,38 +662,96 @@ def build_app():
         with gr.Accordion("Debug: HDF5 tree", open=False):
             inspect_btn = gr.Button("Inspect HDF5 structure")
             hdf5_tree = gr.Textbox(lines=22, label="HDF5 tree")
-            inspect_btn.click(fn=inspect_hdf5_tree, outputs=hdf5_tree)
+            inspect_btn.click(
+                fn=inspect_hdf5_tree,
+                inputs=[preset, custom_repo_id, custom_filename],
+                outputs=hdf5_tree,
+            )
+
+        def update_custom_visibility(preset_name):
+            visible = preset_name == "Custom"
+            return gr.update(visible=visible), gr.update(visible=visible)
+
+        def update_after_dataset_change(preset_name, custom_repo_id, custom_filename):
+            repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
+            n = get_num_trajectories(repo_id, filename)
+            keys = get_available_image_keys(repo_id, filename, 0)
+            traj = load_traj(repo_id, filename, 0)
+            return (
+                gr.update(maximum=max(n - 1, 1), value=0),
+                gr.update(maximum=max(len(traj) - 1, 1), value=0),
+                gr.update(choices=keys, value=keys[:2]),
+                f"Loaded `{repo_id}` / `{filename}`
+Detected trajectories: {n}",
+            )
+
+        dataset_status = gr.Textbox(label="Dataset status", lines=2, value=f"Loaded default dataset
+Detected trajectories: {n_traj}")
+
+        preset.change(
+            fn=update_custom_visibility,
+            inputs=preset,
+            outputs=[custom_repo_id, custom_filename],
+        ).then(
+            fn=update_after_dataset_change,
+            inputs=[preset, custom_repo_id, custom_filename],
+            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],
+            outputs=[gallery, action_plot, info],
+        )
+
+        custom_repo_id.submit(
+            fn=update_after_dataset_change,
+            inputs=[preset, custom_repo_id, custom_filename],
+            outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
+        )
+        custom_filename.submit(
+            fn=update_after_dataset_change,
+            inputs=[preset, custom_repo_id, custom_filename],
+            outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
+        )
 
         traj_slider.change(
             fn=update_after_traj_change,
-            inputs=traj_slider,
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider],
             outputs=[timestep_slider, image_keys],
         ).then(
             fn=render_frame,
-            inputs=[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],
+        )
+
+        # Use release for the timestep slider so the gallery does not clear/re-render
+        # continuously while the user drags through a trajectory.
+        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],
             outputs=[gallery, action_plot, info],
         )
 
-        for widget in [timestep_slider, image_keys, chunk_len, display_scale, reverse_channels]:
+        # These controls can re-render immediately because they are changed less often.
+        for widget in [image_keys, chunk_len, display_scale, reverse_channels]:
             widget.change(
                 fn=render_frame,
-                inputs=[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=[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],
         )
 
         demo.load(
             fn=update_after_traj_change,
-            inputs=traj_slider,
+            inputs=[preset, custom_repo_id, custom_filename, traj_slider],
             outputs=[timestep_slider, image_keys],
         ).then(
             fn=render_frame,
-            inputs=[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],
         )