Auto-display catalog characters after video upload

frame_extraction/src/frame_extraction/app.py

@@ -6,7 +6,7 @@ import os
 import shutil
 import uuid
 from pathlib import Path
-from typing import Any
+from typing import Any, Tuple
 
 import gradio as gr
 import numpy as np
@@ -28,7 +28,30 @@ def ensure_output_dirs() -> None:
     (OUTPUT_DIR / "frames").mkdir(parents=True, exist_ok=True)
 
 
-def build_catalog_from_video(file: gr.FileData) -> tuple[str | None, str]:
+def summarize_catalog(catalog_path: Path) -> Tuple[str, list[dict[str, Any]], list[list[str]]]:
+    if not catalog_path.exists():
+        return ("Catalog not found.", [], [])
+    data = json.loads(catalog_path.read_text(encoding="utf-8"))
+    references = data.get("references", [])
+    message = f"Catalog ready ({len(references)} references)."
+    index = [
+        {
+            "character_id": ref.get("character_id"),
+            "reference_path": ref.get("reference_path"),
+            "frame_path": ref.get("frame_path"),
+            "sharpness": ref.get("sharpness"),
+        }
+        for ref in references
+    ]
+    gallery = [
+        [ref.get("reference_path", ""), ref.get("character_id", "unknown")]
+        for ref in references
+        if ref.get("reference_path")
+    ]
+    return message, index, gallery
+
+
+def build_catalog_from_video(file: gr.FileData) -> tuple[str | None, str, list[dict[str, Any]], list[list[str]]]:
     if file is None:
         raise gr.Error("Please upload a source video first.")
 
@@ -41,10 +64,8 @@ def build_catalog_from_video(file: gr.FileData) -> tuple[str | None, str]:
     catalog_dir = OUTPUT_DIR / "catalogs" / f"catalog_{run_id}"
     cfg = CatalogConfig(video_path=video_path, output_dir=catalog_dir)
     catalog_path = build_catalog(cfg)
-    catalog_data = json.loads(catalog_path.read_text(encoding="utf-8"))
-    ref_count = len(catalog_data.get("references", []))
-    message = f"Catalog ready ({ref_count} references)."
-    return str(catalog_path), message
+    message, index, gallery = summarize_catalog(catalog_path)
+    return str(catalog_path), message, index, gallery
 
 
 def predict_from_arrays(arrays: list[np.ndarray], catalog_path: str | None) -> tuple[list[dict[str, Any]], list[list[str]]]:
@@ -75,15 +96,28 @@ def predict_from_arrays(arrays: list[np.ndarray], catalog_path: str | None) -> t
     gallery_items = [
         [item.get("reference_crop", ""), f"{item.get('character_id', 'unknown')} ({item.get('similarity', 0):.2f})"]
         for item in data
+        if item.get("reference_crop")
     ]
     return data, gallery_items
 
 
+def load_initial_catalog() -> tuple[str | None, str, list[dict[str, Any]], list[list[str]]]:
+    if CATALOG_PATH.exists():
+        message, index, gallery = summarize_catalog(CATALOG_PATH)
+        return str(CATALOG_PATH), message, index, gallery
+    return None, "Upload a video to generate a catalog.", [], []
+
+
 def build_interface() -> gr.Blocks:
+    initial_catalog, initial_status, initial_index, initial_gallery = load_initial_catalog()
+
     with gr.Blocks() as demo:
         gr.Markdown("# Character Reference Matcher")
-        catalog_state = gr.State(str(CATALOG_PATH) if CATALOG_PATH.exists() else None)
-        status_box = gr.Textbox(label="Status", value="Upload a video to generate a catalog.", interactive=False)
+        catalog_state = gr.State(initial_catalog)
+
+        status_box = gr.Textbox(label="Status", value=initial_status, interactive=False)
+        catalog_json = gr.JSON(label="Character Index", value=initial_index)
+        catalog_gallery = gr.Gallery(label="Catalog Characters", columns=4, value=initial_gallery)
 
         video_upload = gr.File(label="Source video", file_types=["video"], height="auto")
         frame_upload = gr.UploadButton(
@@ -93,15 +127,18 @@ def build_interface() -> gr.Blocks:
         )
 
         matches_json = gr.JSON(label="Matches")
-        gallery = gr.Gallery(label="Reference Thumbnails", columns=2)
+        match_gallery = gr.Gallery(label="Matched Characters", columns=3)
+
+        def on_video_upload(file: gr.FileData) -> tuple[str | None, str, list[dict[str, Any]], list[list[str]]]:
+            return build_catalog_from_video(file)
 
-        video_upload.change(build_catalog_from_video, inputs=video_upload, outputs=[catalog_state, status_box])
+        video_upload.change(on_video_upload, inputs=video_upload, outputs=[catalog_state, status_box, catalog_json, catalog_gallery])
 
-        def handle_frames(files: list[gr.FileData], catalog_path: str | None) -> tuple[list[dict[str, Any]], list[list[str]]]:
+        def on_frames_upload(files: list[gr.FileData], catalog_path: str | None) -> tuple[list[dict[str, Any]], list[list[str]]]:
             arrays = [np.array(Image.open(file.name).convert("RGB")) for file in files]
             return predict_from_arrays(arrays, catalog_path)
 
-        frame_upload.upload(handle_frames, inputs=[frame_upload, catalog_state], outputs=[matches_json, gallery])
+        frame_upload.upload(on_frames_upload, inputs=[frame_upload, catalog_state], outputs=[matches_json, match_gallery])
     return demo