Compat with Gradio upload/gallery

frame_extraction/src/frame_extraction/app.py

@@ -7,11 +7,11 @@ from typing import Any
 
 import gradio as gr
 import numpy as np
+from PIL import Image
 
 from .config import MatchConfig
 from .matcher import match_frames
 
-
 CATALOG_PATH = Path(os.getenv("FRAME_CATALOG", "catalog/catalog.json"))
 OUTPUT_DIR = Path(os.getenv("FRAME_OUTPUT_DIR", "app_outputs"))
 
@@ -23,25 +23,28 @@ def load_catalog() -> dict[str, Any] | None:
     return None
 
 
+def ensure_output_dirs() -> None:
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    (OUTPUT_DIR / "inputs").mkdir(parents=True, exist_ok=True)
+
+
 catalog_cache = load_catalog()
 
 
-def predict(image_inputs: list[np.ndarray]) -> tuple[list[dict[str, Any]], list[tuple[str, str]]]:
+def predict_from_arrays(arrays: list[np.ndarray]) -> tuple[list[dict[str, Any]], list[list[str]]]:
     if catalog_cache is None:
         raise gr.Error("Catalog not found. Upload catalog.json or set FRAME_CATALOG.")
 
-    if not image_inputs:
+    if not arrays:
         raise gr.Error("Please upload at least one frame.")
 
+    ensure_output_dirs()
     frames_dir = OUTPUT_DIR / "inputs"
-    frames_dir.mkdir(parents=True, exist_ok=True)
-
-    from PIL import Image
 
     saved_paths: list[Path] = []
-    for idx, image in enumerate(image_inputs):
+    for idx, array in enumerate(arrays):
         output_path = frames_dir / f"upload_{idx:03d}.png"
-        Image.fromarray(image).save(output_path)
+        Image.fromarray(array).save(output_path)
         saved_paths.append(output_path)
 
     output_path = OUTPUT_DIR / "matches.json"
@@ -55,7 +58,8 @@ def predict(image_inputs: list[np.ndarray]) -> tuple[list[dict[str, Any]], list[
     match_frames(cfg)
     data = json.loads(output_path.read_text(encoding="utf-8"))
     gallery_items = [
-        (item["reference_crop"], f"{item['character_id']} ({item['similarity']:.2f})") for item in data
+        [item["reference_crop"], f"{item['character_id']} ({item['similarity']:.2f})"]
+        for item in data
     ]
     return data, gallery_items
 
@@ -63,28 +67,25 @@ def predict(image_inputs: list[np.ndarray]) -> tuple[list[dict[str, Any]], list[
 def build_interface() -> gr.Blocks:
     with gr.Blocks() as demo:
         gr.Markdown("# Character Reference Matcher")
-        image_input = gr.UploadButton(
+        upload = gr.UploadButton(
             label="Upload frames",
             file_types=["image"],
             file_count="multiple",
         )
         matches_json = gr.JSON(label="Matches")
-        gallery = gr.Gallery(label="Reference Thumbnails").style(grid=2)
-
-        def _handle_upload(files: list[gr.FileData]):
-            from PIL import Image
-            import numpy as np
+        gallery = gr.Gallery(label="Reference Thumbnails", columns=2, height="auto")
 
-            images = [np.array(Image.open(file.name)) for file in files]
-            return predict(images)
+        def handle_upload(files: list[gr.FileData]) -> 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)
 
-        image_input.upload(_handle_upload, inputs=image_input, outputs=[matches_json, gallery])
+        upload.upload(handle_upload, inputs=upload, outputs=[matches_json, gallery])
     return demo
 
 
 def main() -> None:
-    interface = build_interface()
-    interface.launch()
+    demo = build_interface()
+    demo.launch()
 
 
 if __name__ == "__main__":