Allow building catalog from uploaded video

frame_extraction/src/frame_extraction/app.py

@@ -1,7 +1,10 @@
+
 from __future__ import annotations
 
 import json
 import os
+import shutil
+import uuid
 from pathlib import Path
 from typing import Any
 
@@ -9,47 +12,60 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
-from .config import MatchConfig
+from .catalog import build_catalog
+from .config import CatalogConfig, MatchConfig
 from .matcher import match_frames
 
-CATALOG_PATH = Path(os.getenv("FRAME_CATALOG", "catalog/catalog.json"))
+CATALOG_ENV_DEFAULT = "catalog/catalog.json"
+CATALOG_PATH = Path(os.getenv("FRAME_CATALOG", CATALOG_ENV_DEFAULT))
 OUTPUT_DIR = Path(os.getenv("FRAME_OUTPUT_DIR", "app_outputs"))
 
 
-def load_catalog() -> dict[str, Any] | None:
-    path = Path(CATALOG_PATH)
-    if path.exists():
-        return json.loads(path.read_text(encoding="utf-8"))
-    return None
-
-
 def ensure_output_dirs() -> None:
     OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-    (OUTPUT_DIR / "inputs").mkdir(parents=True, exist_ok=True)
+    (OUTPUT_DIR / "catalogs").mkdir(parents=True, exist_ok=True)
+    (OUTPUT_DIR / "videos").mkdir(parents=True, exist_ok=True)
+    (OUTPUT_DIR / "frames").mkdir(parents=True, exist_ok=True)
 
 
-catalog_cache = load_catalog()
+def build_catalog_from_video(files: list[gr.FileData]) -> tuple[str | None, str]:
+    if not files:
+        raise gr.Error("Please upload a source video first.")
+
+    ensure_output_dirs()
+    file = files[0]
+    run_id = uuid.uuid4().hex[:8]
+    video_dir = OUTPUT_DIR / "videos"
+    video_path = video_dir / f"{run_id}_{Path(file.name).name}"
+    shutil.copy(file.name, video_path)
+
+    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
 
 
-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.")
+def predict_from_arrays(arrays: list[np.ndarray], catalog_path: str | None) -> tuple[list[dict[str, Any]], list[list[str]]]:
+    if not catalog_path:
+        raise gr.Error("Catalog not ready yet. Upload a video first.")
 
     if not arrays:
         raise gr.Error("Please upload at least one frame.")
 
     ensure_output_dirs()
-    frames_dir = OUTPUT_DIR / "inputs"
+    run_id = uuid.uuid4().hex[:8]
+    frames_dir = OUTPUT_DIR / "frames" / run_id
+    frames_dir.mkdir(parents=True, exist_ok=True)
 
-    saved_paths: list[Path] = []
     for idx, array in enumerate(arrays):
-        output_path = frames_dir / f"upload_{idx:03d}.png"
-        Image.fromarray(array).save(output_path)
-        saved_paths.append(output_path)
+        Image.fromarray(array).save(frames_dir / f"upload_{idx:03d}.png")
 
-    output_path = OUTPUT_DIR / "matches.json"
+    output_path = OUTPUT_DIR / f"matches_{run_id}.json"
     cfg = MatchConfig(
-        catalog_path=Path(CATALOG_PATH),
+        catalog_path=Path(catalog_path),
         frames_dir=frames_dir,
         output_path=output_path,
         top_k=1,
@@ -58,7 +74,7 @@ def predict_from_arrays(arrays: list[np.ndarray]) -> tuple[list[dict[str, Any]],
     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})"]
+        [item.get("reference_crop", ""), f"{item.get('character_id', 'unknown')} ({item.get('similarity', 0):.2f})"]
         for item in data
     ]
     return data, gallery_items
@@ -67,19 +83,31 @@ def predict_from_arrays(arrays: list[np.ndarray]) -> tuple[list[dict[str, Any]],
 def build_interface() -> gr.Blocks:
     with gr.Blocks() as demo:
         gr.Markdown("# Character Reference Matcher")
-        upload = gr.UploadButton(
+        catalog_state = gr.State[str | None](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)
+
+        video_upload = gr.UploadButton(
+            label="Upload source video",
+            file_types=["video"],
+            file_count="single",
+        )
+
+        frame_upload = gr.UploadButton(
             label="Upload frames",
             file_types=["image"],
             file_count="multiple",
         )
+
         matches_json = gr.JSON(label="Matches")
-        gallery = gr.Gallery(label="Reference Thumbnails", columns=2, height="auto")
+        gallery = gr.Gallery(label="Reference Thumbnails", columns=2)
+
+        video_upload.upload(build_catalog_from_video, inputs=video_upload, outputs=[catalog_state, status_box])
 
-        def handle_upload(files: list[gr.FileData]) -> tuple[list[dict[str, Any]], list[list[str]]]:
+        def handle_frames(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)
+            return predict_from_arrays(arrays, catalog_path)
 
-        upload.upload(handle_upload, inputs=upload, outputs=[matches_json, gallery])
+        frame_upload.upload(handle_frames, inputs=[frame_upload, catalog_state], outputs=[matches_json, gallery])
     return demo