Fix

app.py

@@ -1,10 +1,11 @@
+import json
 from functools import lru_cache
 import re
 import traceback
 from typing import Optional
 
 import gradio as gr
-from huggingface_hub import HfApi, hf_hub_url
+from huggingface_hub import HfApi, hf_hub_download, hf_hub_url
 from huggingface_hub.utils import HfHubHTTPError
 
 
@@ -21,6 +22,7 @@ IMAGE_EXTENSIONS = (
 )
 INIT_SCREENSHOT_NAMES = {"intial_screenshot", "initial_screenshot"}
 STEP_FILENAME_PATTERN = re.compile(r"^step_(\d+)(?:\.[^.]+)?$", re.IGNORECASE)
+TRAJECTORY_FILENAME = "traj.jsonl"
 
 
 api = HfApi()
@@ -59,11 +61,72 @@ def _build_path(*parts) -> str:
     return "/".join(part for part in parts if part)
 
 
-def _get_image_urls(repo_id: str, directory: str) -> list[str]:
-    """Return URLs for image files directly within the given directory."""
+def _image_sort_key(path: str):
+    filename = path.rsplit("/", 1)[-1]
+    lower_name = filename.lower()
+
+    if any(lower_name.startswith(name) for name in INIT_SCREENSHOT_NAMES):
+        return (0, 0)
+
+    match = STEP_FILENAME_PATTERN.match(lower_name)
+    if match:
+        return (1, int(match.group(1)))
+
+    return (2, lower_name)
+
+
+def _load_traj_entries(repo_id: str, directory: str) -> list:
+    """Load trajectory annotations from traj.jsonl within the given directory."""
     if not directory:
         return []
 
+    traj_path = _build_path(directory, TRAJECTORY_FILENAME)
+    repo_files = _list_repo_files(repo_id)
+
+    if traj_path not in repo_files:
+        return []
+
+    local_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=traj_path,
+        repo_type="dataset",
+    )
+
+    entries: list = []
+    with open(local_path, "r", encoding="utf-8") as file:
+        for raw_line in file:
+            stripped = raw_line.strip()
+            if not stripped:
+                continue
+
+            parsed = json.loads(stripped)
+            if isinstance(parsed, list):
+                entries.extend(parsed)
+            else:
+                entries.append(parsed)
+
+    return entries
+
+
+def _format_annotation(index: int, annotation) -> str:
+    prefix = f"Step {index + 1}"
+    if isinstance(annotation, str):
+        content = annotation.strip()
+    else:
+        try:
+            content = json.dumps(annotation, ensure_ascii=False)
+        except TypeError:
+            content = str(annotation)
+    return f"{prefix}: {content}" if content else prefix
+
+
+def _prepare_gallery_items(
+    repo_id: str, directory: Optional[str]
+) -> tuple[list[tuple[str, str]], list[str]]:
+    """Prepare gallery items and status messages for a directory."""
+    if not directory:
+        return [], ["Select a fourth-level folder to view screenshots."]
+
     prefix = f"{directory}/"
     files = [path for path in _list_repo_files(repo_id) if path.startswith(prefix)]
 
@@ -75,25 +138,46 @@ def _get_image_urls(repo_id: str, directory: str) -> list[str]:
     ]
 
     sorted_files = sorted(image_files, key=_image_sort_key)
-
-    return [
+    image_urls = [
         hf_hub_url(repo_id=repo_id, filename=path, repo_type="dataset")
         for path in sorted_files
     ]
 
+    status_lines: list[str] = [
+        f"🖼️ Images: {len(image_urls)}",
+    ]
 
-def _image_sort_key(path: str):
-    filename = path.rsplit("/", 1)[-1]
-    lower_name = filename.lower()
+    annotations: list = []
+    try:
+        annotations = _load_traj_entries(repo_id, directory)
+    except json.JSONDecodeError as error:
+        status_lines.append(f"⚠️ Failed to parse `{TRAJECTORY_FILENAME}`: {error}")
+    except Exception as error:  # pragma: no cover - unexpected IO errors
+        status_lines.append(f"⚠️ Error loading `{TRAJECTORY_FILENAME}`: {error}")
 
-    if any(lower_name.startswith(name) for name in INIT_SCREENSHOT_NAMES):
-        return (0, 0)
+    status_lines.append(f"📝 Annotations: {len(annotations)}")
 
-    match = STEP_FILENAME_PATTERN.match(lower_name)
-    if match:
-        return (1, int(match.group(1)))
+    if not image_urls:
+        status_lines.append("⚠️ No images found in this folder.")
 
-    return (2, lower_name)
+    if image_urls and not annotations:
+        status_lines.append(
+            f"⚠️ `{TRAJECTORY_FILENAME}` missing or empty; no annotations to display."
+        )
+
+    if annotations and len(annotations) != len(image_urls):
+        status_lines.append(
+            "⚠️ Mismatch between images and annotations; displaying available pairs."
+        )
+
+    gallery_items: list[tuple[str, str]] = []
+    for idx, url in enumerate(image_urls):
+        caption = "No annotation available"
+        if idx < len(annotations):
+            caption = _format_annotation(idx, annotations[idx])
+        gallery_items.append((url, caption))
+
+    return gallery_items, status_lines
 
 
 def _dropdown_update(
@@ -126,6 +210,7 @@ def refresh_repo(repo_id: str):
             gr.update(choices=[], value=None, interactive=False),
             gr.update(choices=[], value=None, interactive=False),
             gr.update(value=[]),
+            gr.update(value=""),
             gr.update(value=f"❌ Unable to load repo `{repo_id}`: {error}"),
         )
     except Exception as error:  # pragma: no cover - network and auth edge cases
@@ -137,6 +222,7 @@ def refresh_repo(repo_id: str):
             gr.update(choices=[], value=None, interactive=False),
             gr.update(choices=[], value=None, interactive=False),
             gr.update(value=[]),
+            gr.update(value=""),
             gr.update(value=f"❌ Unexpected error loading `{repo_id}`: {error}"),
         )
 
@@ -169,13 +255,12 @@ def refresh_repo(repo_id: str):
     )
     fourth_value = fourth_dirs[0] if fourth_dirs else None
 
-    image_urls = (
-        _get_image_urls(
-            repo_id, _build_path(top_value, second_value, third_value, fourth_value)
-        )
+    target_directory = (
+        _build_path(top_value, second_value, third_value, fourth_value)
         if fourth_value
-        else []
+        else None
     )
+    gallery_items, gallery_status = _prepare_gallery_items(repo_id, target_directory)
 
     first_dropdown_update = _dropdown_update(
         choices=top_dirs,
@@ -209,7 +294,8 @@ def refresh_repo(repo_id: str):
         empty_info="No fourth-level folders under the selection",
     )
 
-    gallery_update = gr.update(value=image_urls)
+    gallery_update = gr.update(value=gallery_items)
+    annotation_update = gr.update(value="\n".join(gallery_status))
 
     return (
         first_dropdown_update,
@@ -217,6 +303,7 @@ def refresh_repo(repo_id: str):
         third_dropdown_update,
         fourth_dropdown_update,
         gallery_update,
+        annotation_update,
         gr.update(value="\n".join(status_lines)),
     )
 
@@ -246,7 +333,10 @@ def update_second_dropdown(repo_id: str, top_level_dir: str):
                 filled_info="Choose a fourth-level folder",
                 empty_info="Select a higher-level folder first",
             )
-            return empty_second, empty_third, empty_fourth, gr.update(value=[])
+            annotation_update = gr.update(
+                value="Select a top-level folder to load screenshots and annotations."
+            )
+            return empty_second, empty_third, empty_fourth, gr.update(value=[]), annotation_update
 
         second_dirs = _get_subdirectories(repo_id, top_level_dir)
         second_value = second_dirs[0] if second_dirs else None
@@ -265,14 +355,12 @@ def update_second_dropdown(repo_id: str, top_level_dir: str):
         )
         fourth_value = fourth_dirs[0] if fourth_dirs else None
 
-        image_urls = (
-            _get_image_urls(
-                repo_id,
-                _build_path(top_level_dir, second_value, third_value, fourth_value),
-            )
+        target_directory = (
+            _build_path(top_level_dir, second_value, third_value, fourth_value)
             if fourth_value
-            else []
+            else None
         )
+        gallery_items, gallery_status = _prepare_gallery_items(repo_id, target_directory)
 
         return (
             _dropdown_update(
@@ -296,7 +384,8 @@ def update_second_dropdown(repo_id: str, top_level_dir: str):
                 filled_info="Choose a fourth-level folder",
                 empty_info="No fourth-level folders under the selection",
             ),
-            gr.update(value=image_urls),
+            gr.update(value=gallery_items),
+            gr.update(value="\n".join(gallery_status)),
         )
     except Exception as error:
         print(f"[update_second_dropdown] Error for {repo_id}/{top_level_dir}: {error}", flush=True)
@@ -322,7 +411,13 @@ def update_second_dropdown(repo_id: str, top_level_dir: str):
             filled_info="Choose a fourth-level folder",
             empty_info="Unable to load subdirectories",
         )
-        return empty_second, empty_third, empty_fourth, gr.update(value=[])
+        return (
+            empty_second,
+            empty_third,
+            empty_fourth,
+            gr.update(value=[]),
+            gr.update(value="Unable to load screenshots or annotations."),
+        )
 
 
 def update_third_dropdown(repo_id: str, top_level_dir: str, second_level_dir: str):
@@ -343,7 +438,10 @@ def update_third_dropdown(repo_id: str, top_level_dir: str, second_level_dir: st
                 filled_info="Choose a fourth-level folder",
                 empty_info="Select higher-level folders first",
             )
-            return empty_third, empty_fourth, gr.update(value=[])
+            annotation_update = gr.update(
+                value="Select higher-level folders to load screenshots and annotations."
+            )
+            return empty_third, empty_fourth, gr.update(value=[]), annotation_update
 
         third_dirs = _get_subdirectories(
             repo_id, _build_path(top_level_dir, second_level_dir)
@@ -359,14 +457,12 @@ def update_third_dropdown(repo_id: str, top_level_dir: str, second_level_dir: st
         )
         fourth_value = fourth_dirs[0] if fourth_dirs else None
 
-        image_urls = (
-            _get_image_urls(
-                repo_id,
-                _build_path(top_level_dir, second_level_dir, third_value, fourth_value),
-            )
+        target_directory = (
+            _build_path(top_level_dir, second_level_dir, third_value, fourth_value)
             if fourth_value
-            else []
+            else None
         )
+        gallery_items, gallery_status = _prepare_gallery_items(repo_id, target_directory)
 
         return (
             _dropdown_update(
@@ -383,7 +479,8 @@ def update_third_dropdown(repo_id: str, top_level_dir: str, second_level_dir: st
                 filled_info="Choose a fourth-level folder",
                 empty_info="No fourth-level folders under the selection",
             ),
-            gr.update(value=image_urls),
+            gr.update(value=gallery_items),
+            gr.update(value="\n".join(gallery_status)),
         )
     except Exception as error:
         print(
@@ -405,7 +502,12 @@ def update_third_dropdown(repo_id: str, top_level_dir: str, second_level_dir: st
             filled_info="Choose a fourth-level folder",
             empty_info="Unable to load subdirectories",
         )
-        return empty_third, empty_fourth, gr.update(value=[])
+        return (
+            empty_third,
+            empty_fourth,
+            gr.update(value=[]),
+            gr.update(value="Unable to load screenshots or annotations."),
+        )
 
 
 def update_fourth_dropdown(
@@ -424,7 +526,10 @@ def update_fourth_dropdown(
                 filled_info="Choose a fourth-level folder",
                 empty_info="Select higher-level folders first",
             )
-            return empty_fourth, gr.update(value=[])
+            annotation_update = gr.update(
+                value="Select higher-level folders to load screenshots and annotations."
+            )
+            return empty_fourth, gr.update(value=[]), annotation_update
 
         fourth_dirs = _get_subdirectories(
             repo_id,
@@ -432,19 +537,17 @@ def update_fourth_dropdown(
         )
         fourth_value = fourth_dirs[0] if fourth_dirs else None
 
-        image_urls = (
-            _get_image_urls(
-                repo_id,
-                _build_path(
-                    top_level_dir,
-                    second_level_dir,
-                    third_level_dir,
-                    fourth_value,
-                ),
+        target_directory = (
+            _build_path(
+                top_level_dir,
+                second_level_dir,
+                third_level_dir,
+                fourth_value,
             )
             if fourth_value
-            else []
+            else None
         )
+        gallery_items, gallery_status = _prepare_gallery_items(repo_id, target_directory)
 
         return (
             _dropdown_update(
@@ -454,7 +557,8 @@ def update_fourth_dropdown(
                 filled_info="Choose a fourth-level folder",
                 empty_info="No fourth-level folders under the selection",
             ),
-            gr.update(value=image_urls),
+            gr.update(value=gallery_items),
+            gr.update(value="\n".join(gallery_status)),
         )
     except Exception as error:
         print(
@@ -470,7 +574,11 @@ def update_fourth_dropdown(
             filled_info="Choose a fourth-level folder",
             empty_info="Unable to load subdirectories",
         )
-        return empty_fourth, gr.update(value=[])
+        return (
+            empty_fourth,
+            gr.update(value=[]),
+            gr.update(value="Unable to load screenshots or annotations."),
+        )
 
 
 def update_gallery(
@@ -483,18 +591,19 @@ def update_gallery(
     """Update the image gallery when the fourth-level selection changes."""
     try:
         if not all([top_level_dir, second_level_dir, third_level_dir, fourth_level_dir]):
-            return gr.update(value=[])
+            annotation_update = gr.update(
+                value="Select all folder levels to load screenshots and annotations."
+            )
+            return gr.update(value=[]), annotation_update
 
-        image_urls = _get_image_urls(
-            repo_id,
-            _build_path(
-                top_level_dir,
-                second_level_dir,
-                third_level_dir,
-                fourth_level_dir,
-            ),
+        directory = _build_path(
+            top_level_dir,
+            second_level_dir,
+            third_level_dir,
+            fourth_level_dir,
         )
-        return gr.update(value=image_urls)
+        gallery_items, gallery_status = _prepare_gallery_items(repo_id, directory)
+        return gr.update(value=gallery_items), gr.update(value="\n".join(gallery_status))
     except Exception as error:
         print(
             "[update_gallery] Error for "
@@ -502,7 +611,7 @@ def update_gallery(
             flush=True,
         )
         print(traceback.format_exc(), flush=True)
-        return gr.update(value=[])
+        return gr.update(value=[]), gr.update(value="Unable to load screenshots or annotations.")
 
 
 with gr.Blocks(title="HF Dataset Explorer") as demo:
@@ -527,6 +636,7 @@ Provide a dataset repository ID (e.g. `org/dataset`) to list its top-level folde
     third_level_dropdown = gr.Dropdown(label="Third-level folders", interactive=False)
     fourth_level_dropdown = gr.Dropdown(label="Fourth-level folders", interactive=False)
     image_gallery = gr.Gallery(label="Images", columns=4)
+    annotation_status = gr.Markdown()
 
     reload_button.click(
         refresh_repo,
@@ -537,6 +647,7 @@ Provide a dataset repository ID (e.g. `org/dataset`) to list its top-level folde
             third_level_dropdown,
             fourth_level_dropdown,
             image_gallery,
+            annotation_status,
             status_display,
         ],
     )
@@ -549,13 +660,19 @@ Provide a dataset repository ID (e.g. `org/dataset`) to list its top-level folde
             third_level_dropdown,
             fourth_level_dropdown,
             image_gallery,
+            annotation_status,
         ],
     )
 
     second_level_dropdown.change(
         update_third_dropdown,
         inputs=[repo_id_input, folder_dropdown, second_level_dropdown],
-        outputs=[third_level_dropdown, fourth_level_dropdown, image_gallery],
+        outputs=[
+            third_level_dropdown,
+            fourth_level_dropdown,
+            image_gallery,
+            annotation_status,
+        ],
     )
 
     third_level_dropdown.change(
@@ -566,7 +683,11 @@ Provide a dataset repository ID (e.g. `org/dataset`) to list its top-level folde
             second_level_dropdown,
             third_level_dropdown,
         ],
-        outputs=[fourth_level_dropdown, image_gallery],
+        outputs=[
+            fourth_level_dropdown,
+            image_gallery,
+            annotation_status,
+        ],
     )
 
     fourth_level_dropdown.change(
@@ -578,7 +699,7 @@ Provide a dataset repository ID (e.g. `org/dataset`) to list its top-level folde
             third_level_dropdown,
             fourth_level_dropdown,
         ],
-        outputs=[image_gallery],
+        outputs=[image_gallery, annotation_status],
     )
 
     demo.load(
@@ -590,6 +711,7 @@ Provide a dataset repository ID (e.g. `org/dataset`) to list its top-level folde
             third_level_dropdown,
             fourth_level_dropdown,
             image_gallery,
+            annotation_status,
             status_display,
         ],
     )