Use upload button for multiple frames

frame_extraction/src/frame_extraction/app.py

@@ -7,7 +7,6 @@ from typing import Any
 
 import gradio as gr
 import numpy as np
-import torch
 
 from .config import MatchConfig
 from .matcher import match_frames
@@ -27,19 +26,23 @@ def load_catalog() -> dict[str, Any] | None:
 catalog_cache = load_catalog()
 
 
-def predict(images: list[np.ndarray]) -> tuple[list[str], list[str]]:
+def predict(image_inputs: list[np.ndarray]) -> tuple[list[dict[str, Any]], list[tuple[str, str]]]:
     if catalog_cache is None:
         raise gr.Error("Catalog not found. Upload catalog.json or set FRAME_CATALOG.")
 
-    if not images:
+    if not image_inputs:
         raise gr.Error("Please upload at least one frame.")
 
     frames_dir = OUTPUT_DIR / "inputs"
     frames_dir.mkdir(parents=True, exist_ok=True)
-    for idx, image in enumerate(images):
-        from PIL import Image
 
-        Image.fromarray(image).save(frames_dir / f"upload_{idx:03d}.png")
+    from PIL import Image
+
+    saved_paths: list[Path] = []
+    for idx, image in enumerate(image_inputs):
+        output_path = frames_dir / f"upload_{idx:03d}.png"
+        Image.fromarray(image).save(output_path)
+        saved_paths.append(output_path)
 
     output_path = OUTPUT_DIR / "matches.json"
     cfg = MatchConfig(
@@ -51,23 +54,37 @@ def predict(images: list[np.ndarray]) -> tuple[list[str], list[str]]:
     )
     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]
+    gallery_items = [
+        (item["reference_crop"], f"{item['character_id']} ({item['similarity']:.2f})") for item in data
+    ]
     return data, gallery_items
 
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Character Reference Matcher")
-    with gr.Row():
-        image_input = gr.Image(type="numpy", image_mode="RGB", label="Upload frames", multiple=True)
-    submit = gr.Button("Match Characters")
-    matches_json = gr.JSON(label="Matches")
-    gallery = gr.Gallery(label="Reference Thumbnails").style(grid=2)
+def build_interface() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        gr.Markdown("# Character Reference Matcher")
+        image_input = 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
+
+            images = [np.array(Image.open(file.name)) for file in files]
+            return predict(images)
 
-    submit.click(predict, inputs=image_input, outputs=[matches_json, gallery])
+        image_input.upload(_handle_upload, inputs=image_input, outputs=[matches_json, gallery])
+    return demo
 
 
 def main() -> None:
-    demo.launch()
+    interface = build_interface()
+    interface.launch()
 
 
 if __name__ == "__main__":