Update app.py

app.py

@@ -46,6 +46,7 @@ DATASET_PRESETS = {
             "20260409_205051_Diffusion_CLIC_intervention_Circular_square_image_abs_"
             "Ta16_offlineFalse_Scale0.01/trajectory_buffer_0.hdf5"
         ),
+        "default_reverse_channels": False,
     },
     "Robosuite Square 20260410": {
         "repo_id": "Zhaoting123/Robosuite_Square_image_abs_with_state",
@@ -53,10 +54,12 @@ DATASET_PRESETS = {
             "20260410_205606_Diffusion_CLIC_intervention_Circular_square_image_abs_"
             "Ta16_offlineFalse_Scale0.01/trajectory_buffer_0.hdf5"
         ),
+        "default_reverse_channels": False,
     },
     "InsertT Nov10 noisy": {
         "repo_id": "Zhaoting123/InsertT",
         "filename": "trajectory_buffer_Nov10_demo_noisy.hdf5",
+        "default_reverse_channels": True,
     },
 }
 
@@ -87,6 +90,20 @@ def resolve_dataset(preset_name, custom_repo_id=None, custom_filename=None):
     return item["repo_id"], item["filename"]
 
 
+def get_default_reverse_channels(preset_name):
+    """Dataset-specific default for BGR<->RGB reversal.
+
+    Robosuite Square presets use normal RGB ordering.
+    InsertT / PushT-style preset requires reversal.
+    Custom datasets default to False so users can still override manually.
+    """
+    preset_name = preset_name or DEFAULT_PRESET
+    if preset_name == "Custom":
+        return False
+    item = DATASET_PRESETS.get(preset_name, DATASET_PRESETS[DEFAULT_PRESET])
+    return bool(item.get("default_reverse_channels", False))
+
+
 @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)
@@ -871,14 +888,18 @@ 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_traj = get_num_trajectories(repo_id, filename)
 
+    reverse_default = get_default_reverse_channels(preset_name)
+
     if n_traj == 0:
         status = "Loaded `{}` / `{}`".format(repo_id, filename)
         status += "\nDetected trajectories: 0"
+        status += "\nreverse_channels default: {}".format(int(reverse_default))
         return (
             gr.update(maximum=1, value=0),
             gr.update(maximum=1, value=0),
             gr.update(choices=[], value=[]),
             status,
+            gr.update(value=reverse_default),
         )
 
     keys = get_available_image_keys(repo_id, filename, 0)
@@ -886,12 +907,14 @@ def update_after_dataset_change(preset_name, custom_repo_id, custom_filename):
 
     status = "Loaded `{}` / `{}`".format(repo_id, filename)
     status += "\nDetected trajectories: {}".format(n_traj)
+    status += "\nreverse_channels default: {}".format(int(reverse_default))
 
     return (
         gr.update(maximum=max(n_traj - 1, 1), value=0),
         gr.update(maximum=max(len(traj) - 1, 1), value=0),
         gr.update(choices=keys, value=keys[:2]),
         status,
+        gr.update(value=reverse_default),
     )
 
 
@@ -992,7 +1015,7 @@ def build_app():
         first_keys = []
         startup_warning = repr(exc)
 
-    default_status = "Loaded default dataset\nDetected trajectories: {}".format(n_traj)
+    default_status = "Loaded default dataset\nDetected trajectories: {}\nreverse_channels default: {}".format(n_traj, int(get_default_reverse_channels(DEFAULT_PRESET)))
 
     with gr.Blocks(title="HDF5 Trajectory Viewer") as demo:
         gr.Markdown(
@@ -1023,7 +1046,7 @@ def build_app():
             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")
-            reverse_channels = gr.Checkbox(value=False, label="Reverse channels BGR↔RGB")
+            reverse_channels = gr.Checkbox(value=get_default_reverse_channels(DEFAULT_PRESET), label="Reverse channels BGR↔RGB")
 
         with gr.Row():
             render_btn = gr.Button("Render frame", variant="primary")
@@ -1062,7 +1085,7 @@ def build_app():
         ).then(
             fn=update_after_dataset_change,
             inputs=[preset, custom_repo_id, custom_filename],
-            outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
+            outputs=[traj_slider, timestep_slider, image_keys, dataset_status, reverse_channels],
         ).then(
             fn=render_frame,
             inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
@@ -1072,12 +1095,12 @@ def build_app():
         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],
+            outputs=[traj_slider, timestep_slider, image_keys, dataset_status, reverse_channels],
         )
         custom_filename.submit(
             fn=update_after_dataset_change,
             inputs=[preset, custom_repo_id, custom_filename],
-            outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
+            outputs=[traj_slider, timestep_slider, image_keys, dataset_status, reverse_channels],
         )
 
         traj_slider.change(
@@ -1130,7 +1153,7 @@ def build_app():
         demo.load(
             fn=update_after_dataset_change,
             inputs=[preset, custom_repo_id, custom_filename],
-            outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
+            outputs=[traj_slider, timestep_slider, image_keys, dataset_status, reverse_channels],
         ).then(
             fn=render_frame,
             inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],