From a1d6759abc0f7015edc58bcec28fd8a4323ad7fc Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Fri, 12 Jun 2026 17:54:23 +0200 Subject: [PATCH] faster calibration curve generation jellyfin intergration --- CMakeLists.txt | 19 ++ README.md | 75 ++++- scripts/filter_gallery.py | 111 +++++++ scripts/make_gallery.py | 98 ++---- scripts/make_jellyfin_gallery.py | 452 ++++++++++++++++++++++++++++ scripts/movienet_eval.py | 48 +-- scripts/run_from_jellyfin.py | 109 +++++++ scripts/sae_embed_loader.py | 36 +++ src/face_embedder_engine.hpp | 104 +++++++ src/gallery/gallery_calibration.hpp | 351 ++++++++++++++++++--- src/gallery/gallery_store.cpp | 26 +- src/gallery/gallery_store.hpp | 4 +- src/nodes/identity_matcher_node.hpp | 14 +- src/nodes/result_sink_node.hpp | 42 ++- src/nodes/scene_tracker_node.hpp | 12 +- src/python_bindings.cpp | 40 +++ src/types.hpp | 4 + 17 files changed, 1379 insertions(+), 166 deletions(-) create mode 100644 scripts/filter_gallery.py create mode 100644 scripts/make_jellyfin_gallery.py create mode 100644 scripts/run_from_jellyfin.py create mode 100644 scripts/sae_embed_loader.py create mode 100644 src/face_embedder_engine.hpp create mode 100644 src/python_bindings.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index a8674df..7ecf078 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -86,6 +86,16 @@ FetchContent_Declare( ) FetchContent_MakeAvailable(nlohmann_json) +# nanobind (Python bindings for the sae_embed module) +find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED) +FetchContent_Declare( + nanobind + GIT_REPOSITORY https://github.com/wjakob/nanobind.git + GIT_TAG v2.4.0 + GIT_SHALLOW TRUE +) +FetchContent_MakeAvailable(nanobind) + # ── Model paths ─────────────────────────────────────────────────────────────── set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models" CACHE PATH "Directory containing ONNX model files") @@ -120,6 +130,15 @@ target_link_libraries(embed_faces PRIVATE ) target_compile_definitions(embed_faces PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}") +# ── sae_embed — Python module: load SCRFD+ArcFace once, embed many images ─── +nanobind_add_module(sae_embed src/python_bindings.cpp) +target_include_directories(sae_embed PRIVATE src) +target_link_libraries(sae_embed PRIVATE + ${OpenCV_LIBS} + onnxruntime +) +target_compile_definitions(sae_embed PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}") + # ── analyze — main analysis binary ─────────────────────────────────────────── add_executable(scene_analyze src/main.cpp) target_link_libraries(scene_analyze PRIVATE sae_gallery) diff --git a/README.md b/README.md index 0cc6c85..c8e7312 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,16 @@ cmake -B build -DCMAKE_BUILD_TYPE=Release cmake --build build -j$(nproc) ``` +This also builds `sae_embed`, a Python module (via nanobind) that loads the +SCRFD detector and ArcFace embedder once and exposes a reusable `embed()` +method. The gallery-builder scripts (`make_gallery.py`, +`make_jellyfin_gallery.py`, `movienet_eval.py`) import it directly — there is +no subprocess fallback, so if it's missing they exit with a build instruction: + +```bash +cmake --build build --target sae_embed +``` + Optional flags: | Flag | Default | Effect | @@ -52,7 +62,7 @@ Models are placed in `external/`: | `scene_analyze_debug` | Same as above + per-frame annotated JPEGs (`SAE_DEBUG=1`) | | `scene_preview` | Live OpenCV display window while analysing | | `build_gallery` | Offline gallery builder from a directory of images | -| `embed_faces` | Standalone embedder used by gallery scripts | +| `sae_embed` | Python module (nanobind) used by gallery-builder scripts — loads SCRFD+ArcFace once | ### `scene_analyze` @@ -74,13 +84,72 @@ Key options: | `--track-max-missing` | — | Frames a track survives without a detection | | `--track-min-frames` | 3 | Observations before a track's mean embedding is used for matching | -### Gallery builder +### Gallery builders + +**Per-movie (TMDB):** ```bash python3 scripts/make_gallery.py --tmdb-bearer --movie-id --output gallery.json ``` -Fetches cast images from TMDB and embeds them via `embed_faces`. +Fetches cast images from TMDB and embeds them via `sae_embed`. + +**Whole-library (Jellyfin):** + +```bash +python3 scripts/make_jellyfin_gallery.py \ + --jellyfin-url http://jellyfin.local:8096 \ + --api-key \ + --output gallery.json +``` + +Scans every Movie/Series in Jellyfin, collects the unique cast across the +whole library, downloads each actor's headshot directly from Jellyfin (no +TMDB key needed), and embeds them via `sae_embed` into one global +gallery.json. Since `identity_matcher` scores faces against the entire +gallery, `scene_analyze` can then recognise any actor from your library in +any film — not just the cast listed for that one title. Pass `--merge` on +later runs to only embed actors newly added to the library. Pass +`--tmdb-key` to fall back to TMDB profile images for actors with no usable +image cached in Jellyfin. + +Jellyfin/TMDB lookups and image downloads for different actors run +concurrently (`--workers`, default 8). Embedding is GPU-bound, so it's +gated separately via `--embed-concurrency` (default 1) — only that many +embed calls run at once while other actors' downloads continue in the +background. + +To restrict a single-title run to that title's credited cast (faster, fewer +look-alike mismatches), filter the global gallery first: + +```bash +python3 scripts/filter_gallery.py \ + --gallery gallery.json \ + --jellyfin-url http://jellyfin.local:8096 \ + --api-key \ + --title "The Matrix" \ + --output gallery_matrix.json +``` + +## Running directly from Jellyfin + +`scripts/run_from_jellyfin.py` resolves a title to its media file via the +Jellyfin API, filters the gallery to that title's cast, and runs +`scene_analyze` in one step. Requires this tool to run on a host that shares +Jellyfin's media mount (it uses the item's on-disk `Path`, not a stream URL): + +```bash +python3 scripts/run_from_jellyfin.py \ + --jellyfin-url http://jellyfin.local:8096 \ + --api-key \ + --title "The Matrix" \ + --gallery gallery.json \ + -- --fps 5 --verbosity 2 +``` + +Anything after `--` is passed through to `scene_analyze` unchanged. Pass +`--no-filter` to use the gallery as-is (skip per-title cast filtering), or +`--item-id` instead of `--title` to skip the search. ## Output format diff --git a/scripts/filter_gallery.py b/scripts/filter_gallery.py new file mode 100644 index 0000000..a516649 --- /dev/null +++ b/scripts/filter_gallery.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""filter_gallery.py — restrict a global gallery.json to one title's known cast. + +make_jellyfin_gallery.py builds a single gallery spanning the whole Jellyfin +library, so identity_matcher can recognise any actor from any film. For a +single-title run you may instead want to restrict matching to that title's +credited cast only — fewer candidates means faster matching and fewer +look-alike false positives. + +This script looks up the title's cast live from Jellyfin and writes a +filtered gallery.json containing only those actors, matched via the +"jellyfin_person_id" field recorded by make_jellyfin_gallery.py. + +Usage: + python scripts/filter_gallery.py \\ + --gallery gallery.json \\ + --jellyfin-url http://jellyfin.local:8096 \\ + --api-key YOUR_API_KEY \\ + --item-id \\ + --output gallery_movie.json + + # Or search by title: + python scripts/filter_gallery.py \\ + --gallery gallery.json \\ + --jellyfin-url http://jellyfin.local:8096 \\ + --api-key YOUR_API_KEY \\ + --title "The Matrix" \\ + --output gallery_movie.json +""" + +import argparse +import json +import sys +from pathlib import Path + +import requests + + +def jf_get(base_url: str, api_key: str, path: str, **params) -> dict: + url = base_url.rstrip("/") + path + headers = {"X-Emby-Token": api_key, "Accept": "application/json"} + r = requests.get(url, params=params, headers=headers, timeout=30) + r.raise_for_status() + return r.json() + + +def find_item_id(base_url: str, api_key: str, title: str, item_types: list[str]) -> str: + data = jf_get( + base_url, api_key, "/Items", + Recursive="true", + IncludeItemTypes=",".join(item_types), + SearchTerm=title, + Limit=10, + ) + items = data.get("Items", []) + if not items: + raise ValueError(f"No item found matching title {title!r}") + if len(items) > 1: + print("Multiple matches found:", file=sys.stderr) + for it in items: + print(f" {it['Id']} {it.get('Type')} {it.get('Name')} ({it.get('ProductionYear')})", file=sys.stderr) + print(f"Using first match: {items[0]['Name']}", file=sys.stderr) + return items[0]["Id"] + + +def fetch_cast_person_ids(base_url: str, api_key: str, item_id: str) -> set[str]: + data = jf_get(base_url, api_key, f"/Items/{item_id}", Fields="People") + return {p["Id"] for p in data.get("People", []) if p.get("Type") == "Actor"} + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--gallery", required=True, + help="Global gallery.json built by make_jellyfin_gallery.py") + parser.add_argument("--jellyfin-url", required=True) + parser.add_argument("--api-key", required=True) + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("--item-id", help="Jellyfin item id of the title") + group.add_argument("--title", help="Title to search for (uses first match)") + parser.add_argument("--item-types", default="Movie,Series", + help="Item types to search when using --title (default: Movie,Series)") + parser.add_argument("--output", required=True, help="Output filtered gallery.json path") + args = parser.parse_args() + + gallery = json.loads(Path(args.gallery).read_text()) + + item_id = args.item_id + if item_id is None: + item_id = find_item_id(args.jellyfin_url, args.api_key, args.title, args.item_types.split(",")) + print(f"Resolved title to item id {item_id}", file=sys.stderr) + + cast_ids = fetch_cast_person_ids(args.jellyfin_url, args.api_key, item_id) + print(f"Title has {len(cast_ids)} credited cast member(s)", file=sys.stderr) + + # current make_jellyfin_gallery.py writes "jellyfin_id"; older galleries used + # "jellyfin_person_id" (see gallery_store.cpp's fallback for the same pair). + actors = [a for a in gallery.get("actors", []) + if (a.get("jellyfin_id") or a.get("jellyfin_person_id")) in cast_ids] + missing = len(cast_ids) - len(actors) + if missing > 0: + print(f"[warn] {missing} cast member(s) not present in gallery (not yet embedded)", file=sys.stderr) + + Path(args.output).write_text(json.dumps({"actors": actors}, indent=2) + "\n") + print(f"Saved {len(actors)} actor(s) to {args.output}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/scripts/make_gallery.py b/scripts/make_gallery.py index 89a6996..fe04554 100755 --- a/scripts/make_gallery.py +++ b/scripts/make_gallery.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 """make_gallery.py — fetch actor images for a movie and build gallery.json. -Fetches the cast from TMDB, downloads actor profile images, runs the C++ -embed_faces binary (SCRFD + ArcFace, same models as scene_analyze) to produce -embeddings, then writes gallery.json. +Fetches the cast from TMDB, downloads actor profile images, embeds them via +the sae_embed module (SCRFD + ArcFace, same models as scene_analyze, loaded +once), then writes gallery.json. Requirements: pip install requests Pillow @@ -22,7 +22,7 @@ Usage: --output gallery.json # Additional options: - # --embed-bin build/embed_faces path to embed_faces binary + # --build-dir build/ build dir containing sae_embed module # --models-dir models/ directory with ONNX models # --max-actors 20 how many cast members to include # --images-per-actor 3 profile images to download per actor @@ -32,19 +32,18 @@ Get a free TMDB API key at: https://www.themoviedb.org/settings/api """ import argparse +import io import json -import os -import subprocess import sys -import tempfile import time from pathlib import Path -import io - import requests from PIL import Image +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from sae_embed_loader import load_embedder + TMDB_BASE = "https://api.themoviedb.org/3" TMDB_IMG = "https://image.tmdb.org/t/p/original" @@ -81,7 +80,7 @@ def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]: # Get IMDB ID for this person ext = tmdb_get(f"/person/{person_id}/external_ids", key) - imdb_id = ext.get("imdb_id") or f"tmdb_{person_id}" + imdb_id = ext.get("imdb_id") or "" # Get profile images (sorted by vote_average desc by TMDB) images_data = tmdb_get(f"/person/{person_id}/images", key) @@ -96,6 +95,7 @@ def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]: "id": person_id, "name": member["name"], "imdb_id": imdb_id, + "tmdb_id": str(person_id), "profile_images": image_urls, }) time.sleep(0.05) # be polite to TMDB @@ -124,42 +124,9 @@ def download_images(actor: dict, dest_dir: Path, n: int) -> list[Path]: return paths -# ── Embedding via embed_faces binary ───────────────────────────────────────── - -def embed_images(image_paths: list[Path], embed_bin: str, - detector: str, arcface: str) -> list[dict | None]: - """ - Call the C++ embed_faces binary on a list of images. - Returns a list of result dicts (or None if no face / error) per image. - """ - if not image_paths: - return [] - - cmd = [ - embed_bin, - "--detector", detector, - "--arcface", arcface, - ] + [str(p) for p in image_paths] - - try: - proc = subprocess.run(cmd, capture_output=True, text=True, check=True) - except subprocess.CalledProcessError as e: - print(f"[error] embed_faces failed:\n{e.stderr}", file=sys.stderr) - return [None] * len(image_paths) - - try: - results = json.loads(proc.stdout) - except json.JSONDecodeError as e: - print(f"[error] embed_faces output is not valid JSON: {e}", file=sys.stderr) - return [None] * len(image_paths) - - return results - - # ── Gallery assembly ───────────────────────────────────────────────────────── -def build_gallery(movie_id: int, key: str, embed_bin: str, - detector: str, arcface: str, +def build_gallery(movie_id: int, key: str, embedder, max_actors: int, images_per_actor: int, image_root: Path) -> dict: """Fetch cast, download images, embed, return gallery dict.""" @@ -171,27 +138,27 @@ def build_gallery(movie_id: int, key: str, embed_bin: str, for actor in actors: safe_name = actor["name"].replace(" ", "_") - actor_dir = image_root / f"{actor['imdb_id']}_{safe_name}" + dir_id = actor["imdb_id"] or f"tmdb_{actor['tmdb_id']}" + actor_dir = image_root / f"{dir_id}_{safe_name}" - print(f"\n{actor['name']} ({actor['imdb_id']})", file=sys.stderr) + print(f"\n{actor['name']} ({dir_id})", file=sys.stderr) image_paths = download_images(actor, actor_dir, images_per_actor) if not image_paths: print(" no images downloaded, skipping", file=sys.stderr) continue print(f" embedding {len(image_paths)} image(s)…", file=sys.stderr) - results = embed_images(image_paths, embed_bin, detector, arcface) embeddings = [] source_images = [] - for path, res in zip(image_paths, results): - if res is None or res.get("embedding") is None: - reason = res.get("error", "unknown") if res else "binary error" - print(f" [skip] {path.name}: {reason}", file=sys.stderr) + for path in image_paths: + res = embedder.embed(str(path)) + if not res.ok: + print(f" [skip] {path.name}: {res.error}", file=sys.stderr) continue - embeddings.append(res["embedding"]) + embeddings.append(res.embedding) source_images.append(path.name) - print(f" [ok] {path.name} conf={res.get('confidence', 0):.2f}", + print(f" [ok] {path.name} conf={res.confidence:.2f}", file=sys.stderr) if not embeddings: @@ -200,6 +167,8 @@ def build_gallery(movie_id: int, key: str, embed_bin: str, gallery_actors.append({ "imdb_id": actor["imdb_id"], + "tmdb_id": actor["tmdb_id"], + "jellyfin_id": "", "name": actor["name"], "source_images": source_images, "embeddings": embeddings, @@ -213,7 +182,7 @@ def build_gallery(movie_id: int, key: str, embed_bin: str, def main(): parser = argparse.ArgumentParser( - description="Fetch TMDB cast images and build gallery.json via embed_faces") + description="Fetch TMDB cast images and build gallery.json via sae_embed") parser.add_argument("--tmdb-key", required=True, help="TMDB Bearer token (API Read Access Token from themoviedb.org/settings/api)") group = parser.add_mutually_exclusive_group(required=True) @@ -222,8 +191,8 @@ def main(): group.add_argument("--movie-id", type=int, help="TMDB movie ID (alternative to --imdb-id)") parser.add_argument("--output", required=True, help="Output gallery.json path") - parser.add_argument("--embed-bin", default="build/embed_faces", - help="Path to embed_faces binary (default: build/embed_faces)") + parser.add_argument("--build-dir", default="build", + help="Build directory containing the sae_embed module (default: build)") parser.add_argument("--models-dir", default="models", help="Directory containing ONNX models (default: models/)") parser.add_argument("--arcface", default=None, @@ -239,21 +208,10 @@ def main(): args = parser.parse_args() # Resolve paths - embed_bin = str(Path(args.embed_bin).resolve()) - models_dir = Path(args.models_dir) - detector = str(models_dir / "scrfd_500m_bnkps.onnx") - arcface = args.arcface if args.arcface else str(models_dir / "arcface_w600k_r50.onnx") output = Path(args.output) image_root = Path(args.image_dir) if args.image_dir else output.parent / "images" - # Validate - if not Path(embed_bin).is_file(): - sys.exit(f"embed_faces binary not found: {embed_bin}\n" - f"Build it first: cmake --build build --target embed_faces") - for model, name in [(detector, "SCRFD"), (arcface, "ArcFace")]: - if not Path(model).is_file(): - sys.exit(f"{name} model not found: {model}\n" - f"Run: bash scripts/download_models.sh") + embedder = load_embedder(args.build_dir, args.models_dir, args.arcface) # Resolve movie ID movie_id = args.movie_id @@ -266,9 +224,7 @@ def main(): gallery = build_gallery( movie_id = movie_id, key = args.tmdb_key, - embed_bin = embed_bin, - detector = detector, - arcface = arcface, + embedder = embedder, max_actors = args.max_actors, images_per_actor = args.images_per_actor, image_root = image_root, diff --git a/scripts/make_jellyfin_gallery.py b/scripts/make_jellyfin_gallery.py new file mode 100644 index 0000000..ac6c8a7 --- /dev/null +++ b/scripts/make_jellyfin_gallery.py @@ -0,0 +1,452 @@ +#!/usr/bin/env python3 +"""make_jellyfin_gallery.py — build a gallery.json spanning an entire Jellyfin library. + +Queries the Jellyfin API for every Movie/Series, collects the unique cast +across the whole library, downloads each actor's headshot directly from +Jellyfin (no TMDB key needed), embeds them with the sae_embed module (SCRFD + +ArcFace, loaded once), and writes one global gallery.json. + +Because identity_matcher scores every detected face against the whole +gallery, scene_analyze can then recognise any actor in your library in any +film — not just the cast TMDB lists for that one title. + +For a single-title run, use scripts/filter_gallery.py afterwards to restrict +matching to that title's credited cast (faster, fewer look-alike mismatches). + +Requirements: + pip install requests Pillow + +Usage: + python scripts/make_jellyfin_gallery.py \\ + --jellyfin-url http://jellyfin.local:8096 \\ + --api-key YOUR_API_KEY \\ + --output gallery.json + + # Re-run later to pick up newly added titles without re-embedding + # actors already in the gallery: + python scripts/make_jellyfin_gallery.py \\ + --jellyfin-url http://jellyfin.local:8096 \\ + --api-key YOUR_API_KEY \\ + --output gallery.json --merge + + # Fall back to TMDB profile images for actors with no usable Jellyfin image: + python scripts/make_jellyfin_gallery.py \\ + --jellyfin-url http://jellyfin.local:8096 \\ + --api-key YOUR_API_KEY \\ + --tmdb-key YOUR_TMDB_KEY \\ + --output gallery.json + +Get a Jellyfin API key from Dashboard → Advanced → API Keys. +Get a free TMDB API key at: https://www.themoviedb.org/settings/api +""" + +import argparse +import concurrent.futures +import io +import json +import sys +from pathlib import Path + +import requests +from PIL import Image + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from sae_embed_loader import load_embedder + +TMDB_BASE = "https://api.themoviedb.org/3" +TMDB_IMG = "https://image.tmdb.org/t/p/original" + + +# ── Jellyfin API helpers ──────────────────────────────────────────────────── + +def jf_get(base_url: str, api_key: str, path: str, **params) -> dict: + url = base_url.rstrip("/") + path + headers = {"X-Emby-Token": api_key, "Accept": "application/json"} + r = requests.get(url, params=params, headers=headers, timeout=30) + r.raise_for_status() + return r.json() + + +def fetch_library_items(base_url: str, api_key: str, item_types: list[str]): + """Yield every Movie/Series item dict (with its People field).""" + start = 0 + limit = 100 + while True: + data = jf_get( + base_url, api_key, "/Items", + Recursive="true", + IncludeItemTypes=",".join(item_types), + Fields="People", + StartIndex=start, + Limit=limit, + ) + items = data.get("Items", []) + for item in items: + yield item + start += limit + if start >= data.get("TotalRecordCount", 0) or not items: + break + + +def collect_actors(base_url: str, api_key: str, item_types: list[str]) -> dict: + """Return {person_id: {name, primary_image_tag, appearances: [...]}}.""" + actors: dict[str, dict] = {} + n_items = 0 + for item in fetch_library_items(base_url, api_key, item_types): + n_items += 1 + for person in item.get("People", []): + if person.get("Type") != "Actor": + continue + pid = person["Id"] + entry = actors.setdefault(pid, { + "name": person["Name"], + "primary_image_tag": person.get("PrimaryImageTag"), + "appearances": [], + }) + if not entry["primary_image_tag"]: + entry["primary_image_tag"] = person.get("PrimaryImageTag") + entry["appearances"].append({ + "item_id": item["Id"], + "title": item.get("Name", ""), + "type": item.get("Type", ""), + }) + if n_items % 50 == 0: + print(f" scanned {n_items} items, {len(actors)} unique actors so far…", file=sys.stderr) + print(f"Scanned {n_items} items, found {len(actors)} unique actors", file=sys.stderr) + return actors + + +def fetch_imdb_id(base_url: str, api_key: str, person_id: str) -> str | None: + try: + data = jf_get(base_url, api_key, f"/Items/{person_id}", Fields="ProviderIds") + except requests.RequestException: + return None + return data.get("ProviderIds", {}).get("Imdb") + + +# ── TMDB fallback (for actors with no usable Jellyfin image) ──────────────── + +def tmdb_get(path: str, key: str, **params) -> dict: + url = TMDB_BASE + path + if key.startswith("eyJ"): + headers = {"Authorization": f"Bearer {key}", "Accept": "application/json"} + r = requests.get(url, params=params, headers=headers, timeout=10) + else: + params["api_key"] = key + r = requests.get(url, params=params, headers={"Accept": "application/json"}, timeout=10) + r.raise_for_status() + return r.json() + + +def tmdb_person_images(tmdb_person_id: str, tmdb_key: str) -> list[str]: + images_data = tmdb_get(f"/person/{tmdb_person_id}/images", tmdb_key) + return [TMDB_IMG + p["file_path"] for p in images_data.get("profiles", []) if p.get("file_path")] + + +def tmdb_person_for_imdb(imdb_id: str, tmdb_key: str) -> tuple[str | None, list[str]]: + """Return (tmdb_person_id, profile_image_urls) for the TMDB person matching this IMDB person id.""" + data = tmdb_get(f"/find/{imdb_id}", tmdb_key, external_source="imdb_id") + people = data.get("person_results", []) + if not people: + return None, [] + tmdb_person_id = str(people[0]["id"]) + return tmdb_person_id, tmdb_person_images(tmdb_person_id, tmdb_key) + + +def tmdb_person_by_name(name: str, tmdb_key: str) -> tuple[str | None, list[str]]: + """Return (tmdb_person_id, profile_image_urls) for the best name match on TMDB. + + Used when Jellyfin has no IMDB ProviderId for this person (the common case — + Jellyfin rarely populates ProviderIds on Person items), so /find/{imdb_id} + isn't an option. /search/person is sorted by popularity; take the top hit. + """ + data = tmdb_get("/search/person", tmdb_key, query=name) + people = data.get("results", []) + if not people: + return None, [] + tmdb_person_id = str(people[0]["id"]) + return tmdb_person_id, tmdb_person_images(tmdb_person_id, tmdb_key) + + +def download_urls(urls: list[str], dest_dir: Path, n: int, start_index: int = 0) -> list[Path]: + """Download up to n images from urls into dest_dir, numbered from start_index.""" + dest_dir.mkdir(parents=True, exist_ok=True) + paths = [] + for i, url in enumerate(urls[:n]): + out = dest_dir / f"{start_index + i:02d}.jpg" + if out.exists() and out.stat().st_size > 1024: + paths.append(out) + continue + try: + r = requests.get(url, timeout=15) + r.raise_for_status() + Image.open(io.BytesIO(r.content)).convert("RGB").save(out, "JPEG") + paths.append(out) + except Exception as e: + print(f" [warn] TMDB image download failed: {url}: {e}", file=sys.stderr) + return paths + + +# ── Image download ────────────────────────────────────────────────────────── + +def download_person_images(base_url: str, api_key: str, person_id: str, + dest_dir: Path, n: int) -> list[Path]: + """Download up to n images for a Jellyfin person item into dest_dir.""" + dest_dir.mkdir(parents=True, exist_ok=True) + headers = {"X-Emby-Token": api_key} + paths = [] + for i in range(n): + out = dest_dir / f"{i:02d}.jpg" + if out.exists() and out.stat().st_size > 1024: + paths.append(out) + continue + url = base_url.rstrip("/") + f"/Items/{person_id}/Images/Primary/{i}" + try: + r = requests.get(url, headers=headers, params={"api_key": api_key}, timeout=15) + if r.status_code == 404: + break + r.raise_for_status() + ctype = r.headers.get("Content-Type", "") + if not ctype.startswith("image/"): + print(f" [warn] unexpected response for {person_id} index {i}: " + f"status={r.status_code} content-type={ctype!r} " + f"len={len(r.content)} body={r.content[:200]!r}", file=sys.stderr) + break + Image.open(io.BytesIO(r.content)).convert("RGB").save(out, "JPEG") + paths.append(out) + except Exception as e: + print(f" [warn] image download failed for {person_id} index {i}: {e}", file=sys.stderr) + break + return paths + + +# ── Per-actor pipeline ─────────────────────────────────────────────────────── + +def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict, + images_per_actor: int, actor_dir: Path, + fetch_imdb_ids: bool, tmdb_key: str | None + ) -> tuple[list[Path], str | None, str | None]: + """Network-bound: download Jellyfin image(s), then fall back to TMDB if short.""" + name = info["name"] + print(f"{name} ({pid}) — in {len(info['appearances'])} title(s)", file=sys.stderr) + image_paths = download_person_images(base_url, api_key, pid, actor_dir, images_per_actor) + + imdb_id = None + if fetch_imdb_ids or tmdb_key: + imdb_id = fetch_imdb_id(base_url, api_key, pid) + + # Resolve the TMDB person id whenever possible — independent of whether + # Jellyfin already gave us enough images, so tmdb_id is always populated + # when --tmdb-key is set. Jellyfin Person items rarely have an IMDB + # ProviderId, so fall back to a name search when imdb_id is unknown. + tmdb_id = None + tmdb_urls: list[str] = [] + if tmdb_key: + try: + if imdb_id: + tmdb_id, tmdb_urls = tmdb_person_for_imdb(imdb_id, tmdb_key) + if tmdb_id is None: + tmdb_id, tmdb_urls = tmdb_person_by_name(name, tmdb_key) + except requests.RequestException as e: + print(f" [warn] {name}: TMDB lookup failed: {e}", file=sys.stderr) + + if len(image_paths) < images_per_actor and tmdb_urls: + needed = images_per_actor - len(image_paths) + print(f" {name}: Jellyfin image missing/incomplete, falling back to TMDB " + f"({len(tmdb_urls)} image(s) available)…", file=sys.stderr) + image_paths += download_urls(tmdb_urls, actor_dir, needed, start_index=len(image_paths)) + + return image_paths, imdb_id, tmdb_id + + +def embed_actor(pid: str, info: dict, image_paths: list[Path], + imdb_id: str | None, tmdb_id: str | None, embedder, fetch_imdb_ids: bool, + embed_executor: concurrent.futures.ThreadPoolExecutor) -> dict | None: + """GPU-bound: run sae_embed, always on embed_executor's single dedicated thread. + + onnxruntime's CUDA EP / cudnn_frontend execution plans are not safe to run + from arbitrary threads — calling Run() from a different OS thread than the + one that last used the session corrupts the cudnn graph (CUDNN_FE failure + 11 / CUDNN_BACKEND_API_FAILED). Routing every embed() call through one + persistent thread avoids that regardless of how many worker threads are + fetching images concurrently. + """ + name = info["name"] + if not image_paths: + print(f" {name}: no image available, skipping", file=sys.stderr) + return None + + embeddings = [] + source_images = [] + print(f" {name}: embedding {len(image_paths)} image(s)…", file=sys.stderr) + for path in image_paths: + res = embed_executor.submit(embedder.embed, str(path)).result() + if not res.ok: + print(f" [skip] {name}/{path.name}: {res.error}", file=sys.stderr) + continue + embeddings.append(res.embedding) + source_images.append(path.name) + + if not embeddings: + print(f" {name}: no valid embeddings, skipping actor", file=sys.stderr) + return None + + print(f" {name}: → {len(embeddings)} embedding(s) stored", file=sys.stderr) + return { + "imdb_id": imdb_id if (fetch_imdb_ids and imdb_id) else "", + "tmdb_id": tmdb_id or "", + "jellyfin_id": pid, + "name": name, + "source_images": source_images, + "embeddings": embeddings, + "appearances": info["appearances"], + } + + +def process_actor(pid: str, info: dict, base_url: str, api_key: str, + embedder, images_per_actor: int, image_root: Path, + fetch_imdb_ids: bool, tmdb_key: str | None, + embed_executor: concurrent.futures.ThreadPoolExecutor) -> dict | None: + safe_name = info["name"].replace(" ", "_") + actor_dir = image_root / f"{pid}_{safe_name}" + image_paths, imdb_id, tmdb_id = fetch_actor_images( + base_url, api_key, pid, info, images_per_actor, actor_dir, fetch_imdb_ids, tmdb_key) + return embed_actor(pid, info, image_paths, imdb_id, tmdb_id, embedder, fetch_imdb_ids, embed_executor) + + +# ── Gallery assembly ───────────────────────────────────────────────────────── + +def build_gallery(base_url: str, api_key: str, embedder, item_types: list[str], + images_per_actor: int, image_root: Path, + fetch_imdb_ids: bool, existing_actors: dict, + tmdb_key: str | None = None, workers: int = 8) -> dict: + actors = collect_actors(base_url, api_key, item_types) + + gallery_actors = [] + todo = [] + for pid, info in actors.items(): + if pid in existing_actors: + gallery_actors.append(existing_actors[pid]) + else: + todo.append((pid, info)) + + if len(gallery_actors): + print(f"Skipping {len(gallery_actors)} actor(s) already present in existing gallery", file=sys.stderr) + print(f"Processing {len(todo)} new actor(s) with {workers} worker(s)…", file=sys.stderr) + + n_done = 0 + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor, \ + concurrent.futures.ThreadPoolExecutor(max_workers=1, thread_name_prefix="embed") as embed_executor: + futures = { + executor.submit(process_actor, pid, info, base_url, api_key, embedder, + images_per_actor, image_root, + fetch_imdb_ids, tmdb_key, embed_executor): info["name"] + for pid, info in todo + } + for future in concurrent.futures.as_completed(futures): + n_done += 1 + name = futures[future] + try: + actor = future.result() + except Exception as e: + print(f" [error] {name}: {e}", file=sys.stderr) + continue + if actor: + gallery_actors.append(actor) + if n_done % 25 == 0: + print(f" progress: {n_done}/{len(todo)} actors processed", file=sys.stderr) + + return {"actors": gallery_actors} + + +# ── Entry point ─────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser( + description="Build a gallery.json spanning an entire Jellyfin library", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--jellyfin-url", required=True, + help="Jellyfin base server URL only, e.g. http://jellyfin.local:8096 " + "(no /Items or other API path)") + parser.add_argument("--api-key", required=True, + help="Jellyfin API key (Dashboard → Advanced → API Keys)") + parser.add_argument("--output", required=True, help="Output gallery.json path") + parser.add_argument("--item-types", default="Movie,Series", + help="Comma-separated Jellyfin item types to scan (default: Movie,Series)") + parser.add_argument("--build-dir", default="build", + help="Build directory containing the sae_embed module (default: build)") + parser.add_argument("--models-dir", default="models", + help="Directory containing ONNX models (default: models/)") + parser.add_argument("--arcface", default=None, + help="Path to ArcFace ONNX model (overrides --models-dir selection)") + parser.add_argument("--images-per-actor", type=int, default=1, + help="Images to download per actor (default: 1 — Jellyfin usually caches one)") + parser.add_argument("--image-dir", default=None, + help="Where to store downloaded images (default: /images)") + parser.add_argument("--fetch-imdb-ids", action="store_true", + help="Resolve each actor's real IMDB id via Jellyfin ProviderIds " + "(one extra API call per new actor; otherwise imdb_id is left empty)") + parser.add_argument("--tmdb-key", default=None, + help="TMDB API key/bearer token. If set, actors with no usable " + "Jellyfin image fall back to TMDB profile images (looked up " + "via the actor's IMDB id, requires one extra Jellyfin call per actor)") + parser.add_argument("--merge", action="store_true", + help="If --output already exists, keep its actors and only embed " + "actors not already present (matched by jellyfin_id)") + parser.add_argument("--workers", type=int, default=8, + help="Concurrent worker threads for Jellyfin/TMDB lookups and " + "image downloads (default: 8). Embedding itself always runs " + "on a single dedicated thread regardless of this value.") + args = parser.parse_args() + + # Defend against a base URL that accidentally includes an API path, + # e.g. "https://host/Items?" — strip any query string and trailing + # /Items so we don't build double-nested, query-mangled URLs. + jellyfin_url = args.jellyfin_url.split("?", 1)[0].rstrip("/") + if jellyfin_url.endswith("/Items"): + jellyfin_url = jellyfin_url[: -len("/Items")] + + output = Path(args.output) + image_root = Path(args.image_dir) if args.image_dir else output.parent / "images" + item_types = [t.strip() for t in args.item_types.split(",") if t.strip()] + + embedder = load_embedder(args.build_dir, args.models_dir, args.arcface) + + existing_actors = {} + if args.merge and output.is_file(): + existing = json.loads(output.read_text()) + for actor in existing.get("actors", []): + # older galleries used "jellyfin_person_id" + pid = actor.get("jellyfin_id") or actor.get("jellyfin_person_id") + if pid: + existing_actors[pid] = actor + print(f"Loaded {len(existing_actors)} actor(s) from existing gallery for merge", file=sys.stderr) + + gallery = build_gallery( + base_url=jellyfin_url, + api_key=args.api_key, + embedder=embedder, + item_types=item_types, + images_per_actor=args.images_per_actor, + image_root=image_root, + fetch_imdb_ids=args.fetch_imdb_ids, + existing_actors=existing_actors, + tmdb_key=args.tmdb_key, + workers=args.workers, + ) + + n_actors = len(gallery["actors"]) + n_embeddings = sum(len(a["embeddings"]) for a in gallery["actors"]) + print(f"\nGallery: {n_actors} actors, {n_embeddings} total embeddings", file=sys.stderr) + + if n_actors == 0: + sys.exit("No actors could be processed — check Jellyfin URL/API key and models.") + + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(gallery, indent=2) + "\n") + print(f"Saved: {output}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/scripts/movienet_eval.py b/scripts/movienet_eval.py index 3084c48..4a40b6c 100644 --- a/scripts/movienet_eval.py +++ b/scripts/movienet_eval.py @@ -8,8 +8,7 @@ Usage: --arcface models/arcface_w600k_r50.onnx \ --gt eval/gt.json \ --output eval/predictions_r50.json \ - [--yunet models/face_detection_yunet_2023mar.onnx] \ - [--embed-bin build/embed_faces] + [--build-dir build] Input (--gt): list of {"crop": , "imdb_id": , "actor_name": } Output: list of {"crop", "gt", "pred", "similarity", "detection_failed", "all_scores"} @@ -17,11 +16,12 @@ Output: list of {"crop", "gt", "pred", "similarity", "detection_failed", "all_sc import argparse import json -import math -import subprocess import sys from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from sae_embed_loader import load_embedder + def load_gallery(path: str) -> dict[str, dict]: """Return {imdb_id: {"name": str, "embeddings": [[float]]}}.""" @@ -35,22 +35,6 @@ def dot(a: list[float], b: list[float]) -> float: return sum(x * y for x, y in zip(a, b)) -def embed_images(paths: list[Path], embed_bin: str, yunet: str, arcface: str) -> list[dict | None]: - if not paths: - return [] - cmd = [embed_bin, "--yunet", yunet, "--arcface", arcface] + [str(p) for p in paths] - try: - proc = subprocess.run(cmd, capture_output=True, text=True, check=True) - except subprocess.CalledProcessError as e: - print(f"[error] embed_faces failed:\n{e.stderr}", file=sys.stderr) - return [None] * len(paths) - try: - return json.loads(proc.stdout) - except json.JSONDecodeError as e: - print(f"[error] embed_faces JSON parse error: {e}", file=sys.stderr) - return [None] * len(paths) - - def match(embedding: list[float], gallery: dict[str, dict]) -> tuple[str, float, dict[str, float]]: """Return (best_imdb_id, best_similarity, {imdb_id: similarity}).""" scores: dict[str, float] = {} @@ -67,10 +51,14 @@ def main(): p.add_argument("--arcface", required=True) p.add_argument("--gt", required=True) p.add_argument("--output", required=True) - p.add_argument("--yunet", default="models/face_detection_yunet_2023mar.onnx") - p.add_argument("--embed-bin", default="build/embed_faces") + p.add_argument("--build-dir", default="build", + help="Build directory containing the sae_embed module (default: build)") + p.add_argument("--models-dir", default="models", + help="Directory containing ONNX models (default: models/)") args = p.parse_args() + embedder = load_embedder(args.build_dir, args.models_dir, args.arcface) + gallery = load_gallery(args.gallery) print(f"[eval] gallery: {len(gallery)} actors", file=sys.stderr) @@ -78,29 +66,19 @@ def main(): gt_entries = json.load(f) print(f"[eval] probe crops: {len(gt_entries)}", file=sys.stderr) - # Batch all crops in one embed_faces call to amortise startup cost crop_paths = [Path(e["crop"]) for e in gt_entries] missing = [p for p in crop_paths if not p.exists()] if missing: print(f"[warn] {len(missing)} crop(s) not found on disk, skipping", file=sys.stderr) - results_raw = embed_images( - [p for p in crop_paths if p.exists()], - args.embed_bin, args.yunet, args.arcface - ) - - # Re-index results back to original list (missing files get None) - raw_iter = iter(results_raw) - embed_results: list[dict | None] = [] - for p in crop_paths: - embed_results.append(next(raw_iter) if p.exists() else None) + embed_results = [embedder.embed(str(p)) if p.exists() else None for p in crop_paths] predictions = [] n_det_fail = 0 n_correct = 0 for entry, result in zip(gt_entries, embed_results): - detection_failed = result is None or result.get("embedding") is None + detection_failed = result is None or not result.ok if detection_failed: n_det_fail += 1 predictions.append({ @@ -113,7 +91,7 @@ def main(): }) continue - pred_id, sim, all_scores = match(result["embedding"], gallery) + pred_id, sim, all_scores = match(result.embedding, gallery) correct = pred_id == entry["imdb_id"] if correct: n_correct += 1 diff --git a/scripts/run_from_jellyfin.py b/scripts/run_from_jellyfin.py new file mode 100644 index 0000000..b1e7bfc --- /dev/null +++ b/scripts/run_from_jellyfin.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""run_from_jellyfin.py — resolve a Jellyfin title to its media file and run scene_analyze. + +Looks up a Movie/Episode in Jellyfin, reads its on-disk Path (Jellyfin and this +tool must share the same media mount), filters the gallery down to that +title's credited cast (via filter_gallery's logic, fewer look-alike +mismatches), and runs scene_analyze against the resolved file. + +Usage: + python scripts/run_from_jellyfin.py \\ + --jellyfin-url http://jellyfin.local:8096 \\ + --api-key YOUR_API_KEY \\ + --title "The Matrix" \\ + --gallery whole_gallery.json \\ + -- --fps 5 --verbosity 2 + +Anything after "--" is passed through unchanged to scene_analyze. +""" + +import argparse +import json +import subprocess +import sys +import tempfile +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from filter_gallery import jf_get, find_item_id, fetch_cast_person_ids + + +def fetch_item_path(base_url: str, api_key: str, item_id: str) -> tuple[str, str]: + """Return (Name, Path) for a Jellyfin item.""" + data = jf_get(base_url, api_key, f"/Items/{item_id}", Fields="Path") + path = data.get("Path") + if not path: + raise ValueError(f"Item {item_id} has no Path (not a single media file?)") + return data.get("Name", item_id), path + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--jellyfin-url", required=True) + parser.add_argument("--api-key", required=True) + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("--item-id", help="Jellyfin item id of the title") + group.add_argument("--title", help="Title to search for (uses first match)") + parser.add_argument("--item-types", default="Movie,Episode", + help="Item types to search when using --title (default: Movie,Episode)") + parser.add_argument("--gallery", required=True, + help="Global gallery.json built by make_jellyfin_gallery.py") + parser.add_argument("--no-filter", action="store_true", + help="Skip per-title cast filtering and pass --gallery through as-is") + parser.add_argument("--output", default=None, + help="scene_analyze output JSON (default: .json)") + parser.add_argument("--bin", default="build/scene_analyze", + help="Path to scene_analyze binary (default: build/scene_analyze)") + parser.add_argument("--dry-run", action="store_true", + help="Resolve and print the scene_analyze command without running it") + args, extra = parser.parse_known_args() + if extra and extra[0] == "--": + extra = extra[1:] + + item_id = args.item_id + if item_id is None: + item_id = find_item_id(args.jellyfin_url, args.api_key, args.title, args.item_types.split(",")) + print(f"Resolved title to item id {item_id}", file=sys.stderr) + + name, movie_path = fetch_item_path(args.jellyfin_url, args.api_key, item_id) + print(f"Resolved {name!r} -> {movie_path}", file=sys.stderr) + if not Path(movie_path).is_file(): + sys.exit(f"Resolved path does not exist on this filesystem: {movie_path}\n" + f"(this tool must share Jellyfin's media mount)") + + output = args.output or f"{name}.json" + + gallery_path = args.gallery + filtered_file = None + if not args.no_filter: + cast_ids = fetch_cast_person_ids(args.jellyfin_url, args.api_key, item_id) + gallery = json.loads(Path(args.gallery).read_text()) + actors = [a for a in gallery.get("actors", []) + if (a.get("jellyfin_id") or a.get("jellyfin_person_id")) in cast_ids] + print(f"Filtered gallery to {len(actors)}/{len(gallery.get('actors', []))} " + f"actor(s) credited in {name!r}", file=sys.stderr) + filtered_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", prefix="sae_gallery_", delete=False) + json.dump({"actors": actors}, filtered_file) + filtered_file.close() + gallery_path = filtered_file.name + + cmd = [args.bin, "--movie", movie_path, "--gallery", gallery_path, + "--output", output, *extra] + print("Running:", " ".join(cmd), file=sys.stderr) + + if args.dry_run: + return + + try: + subprocess.run(cmd, check=True) + finally: + if filtered_file is not None: + Path(filtered_file.name).unlink(missing_ok=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/sae_embed_loader.py b/scripts/sae_embed_loader.py new file mode 100644 index 0000000..364be83 --- /dev/null +++ b/scripts/sae_embed_loader.py @@ -0,0 +1,36 @@ +"""Shared loader for the sae_embed nanobind module (SCRFD + ArcFace). + +sae_embed.FaceEmbedder loads both ONNX sessions once and exposes an +embed(path) -> FaceResult method, avoiding the per-process model reload cost +of spawning the embed_faces CLI binary for every image. +""" + +import sys +from pathlib import Path + + +def load_embedder(build_dir: str, models_dir: str, arcface: str | None = None, + conf: float = 0.5, nms: float = 0.4, max_side: int = 500): + """Import sae_embed from build_dir and construct a FaceEmbedder. + + Exits with a clear error if the module or models are missing — there is + no subprocess fallback. + """ + build_path = Path(build_dir).resolve() + sys.path.insert(0, str(build_path)) + try: + import sae_embed + except ImportError as e: + sys.exit( + f"sae_embed module not found in {build_path}: {e}\n" + f"Build it first: cmake --build {build_dir} --target sae_embed" + ) + + models_path = Path(models_dir) + detector_path = str(models_path / "scrfd_500m_bnkps.onnx") + arcface_path = arcface if arcface else str(models_path / "arcface_w600k_r50.onnx") + for model, name in [(detector_path, "SCRFD"), (arcface_path, "ArcFace")]: + if not Path(model).is_file(): + sys.exit(f"{name} model not found: {model}\nRun: bash scripts/download_models.sh") + + return sae_embed.FaceEmbedder(detector_path, arcface_path, conf, nms, max_side) diff --git a/src/face_embedder_engine.hpp b/src/face_embedder_engine.hpp new file mode 100644 index 0000000..ab831ae --- /dev/null +++ b/src/face_embedder_engine.hpp @@ -0,0 +1,104 @@ +#pragma once +// FaceEmbedderEngine — load SCRFD + ArcFace once, embed many images. +// +// Extracted from embed_faces.cpp so the same detect→align→embed pipeline can +// be driven from a long-lived process (the sae_embed Python module) instead +// of a fresh CLI invocation per image, which would reload both ONNX sessions +// every time. + +#include "arcface_embedder.hpp" +#include "face_utils.hpp" +#include "ort_provider.hpp" +#include "scrfd_decoder.hpp" +#include "types.hpp" + +#include <opencv2/imgcodecs.hpp> +#include <opencv2/imgproc.hpp> + +#include <algorithm> +#include <array> +#include <iostream> +#include <memory> +#include <string> + +struct FaceEmbedResult { + bool ok{false}; + std::string error; + Embedding embedding{}; + float confidence{0.f}; + float bbox[4]{}; // x, y, w, h + std::array<cv::Point2f, 5> landmarks{}; +}; + +class FaceEmbedderEngine { +public: + FaceEmbedderEngine(const std::string& detector_model, + const std::string& arcface_model, + float conf = 0.5f, float nms = 0.4f, int max_side = 500) + : max_side_(max_side) + { + const OrtProvider provider = detect_ort_provider(); + std::cerr << "[FaceEmbedderEngine] inference provider: " + << provider_name(provider) << "\n"; + detector_ = std::make_unique<SCRFDDecoder>(detector_model, conf, nms, provider); + embedder_ = std::make_unique<ArcFaceEmbedder>(arcface_model, provider); + } + + FaceEmbedResult embed_path(const std::string& path) const { + cv::Mat img = cv::imread(path); + if (img.empty()) { + FaceEmbedResult res; + res.error = "cannot read image"; + return res; + } + return embed_mat(img); + } + + FaceEmbedResult embed_mat(cv::Mat img) const { + FaceEmbedResult res; + + if (max_side_ > 0) { + const int big = std::max(img.cols, img.rows); + if (big > max_side_) { + const double s = static_cast<double>(max_side_) / big; + cv::resize(img, img, {}, s, s, cv::INTER_AREA); + } + } + + std::vector<DetectedFace> faces = detector_->detect(img); + if (faces.empty()) { + res.error = "no face detected"; + return res; + } + if (faces.size() > 1) + std::cerr << "[warn] " << faces.size() + << " faces detected, using highest-confidence one\n"; + + const auto& best = *std::max_element( + faces.begin(), faces.end(), + [](const DetectedFace& a, const DetectedFace& b) { + return a.confidence < b.confidence; + }); + + cv::Mat crop = align_face(img, best.landmarks); + if (crop.empty()) { + res.error = "alignment failed"; + return res; + } + + res.ok = true; + res.embedding = embedder_->embed_one(crop); + res.confidence = best.confidence; + res.bbox[0] = best.bbox.x; + res.bbox[1] = best.bbox.y; + res.bbox[2] = best.bbox.width; + res.bbox[3] = best.bbox.height; + res.landmarks = best.landmarks; + return res; + } + +private: + std::unique_ptr<SCRFDDecoder> detector_; + std::unique_ptr<ArcFaceEmbedder> embedder_; + int max_side_; +}; diff --git a/src/gallery/gallery_calibration.hpp b/src/gallery/gallery_calibration.hpp index 906855a..5784272 100644 --- a/src/gallery/gallery_calibration.hpp +++ b/src/gallery/gallery_calibration.hpp @@ -1,10 +1,24 @@ #pragma once #include "types.hpp" +#include <algorithm> +#include <chrono> #include <cmath> +#include <cstdint> +#include <cstdlib> +#include <cstring> +#include <fstream> #include <iostream> +#include <stdexcept> +#include <string> +#include <string_view> #include <vector> +#include <nlohmann/json.hpp> +#include <opencv2/core/ocl.hpp> +#include <opencv2/imgcodecs.hpp> +#include <opencv2/imgproc.hpp> + // ── GalleryCalibration ──────────────────────────────────────────────────────── // Platt-style sigmoid calibration: P(match) = σ(a · similarity + b) // where similarity = cosine similarity ∈ [-1, 1] (dot product of L2-normalised @@ -43,24 +57,147 @@ inline GalleryCalibration calibrate_gallery( const std::vector<Embedding>& flat_emb, const std::vector<int>& flat_actor) { - const int n = static_cast<int>(flat_emb.size()); + constexpr int kMinEmbeddingsForPositive = 5; + constexpr float kDedupSimThreshold = 1.f - 1e-7f; // sim above this = duplicate + constexpr int kHistBins = 200; - std::vector<float> X; // cosine similarities - std::vector<float> Y; // labels: 1 = same actor, 0 = different + // ── Per-actor dedup, then eligibility filter ──────────────────────────── + // Drop near-identical duplicate embeddings within each actor (e.g. the + // same source image embedded twice). Actors left with fewer than + // kMinEmbeddingsForPositive distinct embeddings can't supply meaningful + // same-actor (positive) pairs, but still contribute negative + // (cross-actor) pairs for the p(unknown) side of the fit. + int n_actors = 0; + for (int a : flat_actor) n_actors = std::max(n_actors, a + 1); - for (int i = 0; i < n; ++i) { - for (int j = i + 1; j < n; ++j) { - float dot = 0.f; - for (int k = 0; k < 512; ++k) - dot += flat_emb[i][k] * flat_emb[j][k]; - X.push_back(dot); - Y.push_back(flat_actor[i] == flat_actor[j] ? 1.f : 0.f); + std::vector<std::vector<Embedding>> by_actor(n_actors); + for (int i = 0; i < static_cast<int>(flat_emb.size()); ++i) + by_actor[flat_actor[i]].push_back(flat_emb[i]); + + std::vector<Embedding> flat_emb_dedup; + std::vector<int> flat_actor_dedup; + std::vector<bool> actor_eligible(n_actors, false); + int n_eligible = 0; + + for (int ai = 0; ai < n_actors; ++ai) { + std::vector<Embedding> kept; + for (const auto& e : by_actor[ai]) { + bool dup = false; + for (const auto& k : kept) { + if (cosine_similarity(e, k) > kDedupSimThreshold) { dup = true; break; } + } + if (!dup) kept.push_back(e); + } + if (static_cast<int>(kept.size()) >= kMinEmbeddingsForPositive) { + actor_eligible[ai] = true; + ++n_eligible; + } + for (auto& e : kept) { + flat_emb_dedup.push_back(e); + flat_actor_dedup.push_back(ai); } } - int n_pos = 0; - for (float y : Y) if (y > 0.5f) ++n_pos; - int n_neg = static_cast<int>(Y.size()) - n_pos; + const int n = static_cast<int>(flat_emb_dedup.size()); + std::cerr << "[calibration] dedup: " << flat_emb.size() << " -> " << n + << " embeddings (" << n_eligible << "/" << n_actors + << " actors have >= " << kMinEmbeddingsForPositive + << " distinct embeddings, eligible for positive pairs)\n"; + + // Pairwise cosine-similarity matrix S = E * E^T, computed via cv::gemm + // (BLAS/SIMD on CPU) instead of a naive O(n^2 * 512) scalar loop. + // + // SAE_CALIB_GEMM controls the backend: + // "cpu" - cv::Mat gemm only (default if OpenCL unavailable) + // "gpu" - cv::UMat gemm via OpenCL, falls back to CPU if unavailable + // "bench" - run both and report timings + max abs diff (default) + cv::Mat E(n, 512, CV_32F); + for (int i = 0; i < n; ++i) + std::memcpy(E.ptr<float>(i), flat_emb_dedup[i].data(), 512 * sizeof(float)); + + std::string mode = [] { + const char* env = std::getenv("SAE_CALIB_GEMM"); + return env ? std::string(env) : std::string("bench"); + }(); + + bool have_ocl = cv::ocl::haveOpenCL(); + std::cerr << "[calibration] OpenCL available: " << (have_ocl ? "yes" : "no"); + if (have_ocl) { + auto dev = cv::ocl::Device::getDefault(); + std::cerr << " device=\"" << dev.name() << "\"" + << " type=" << (dev.type() == cv::ocl::Device::TYPE_GPU ? "GPU" + : dev.type() == cv::ocl::Device::TYPE_CPU ? "CPU" + : "OTHER"); + } + std::cerr << " (SAE_CALIB_GEMM=" << mode << ")\n"; + + std::cerr << "[calibration] similarity matrix: " << n << "x512 (" + << (n * 512LL * sizeof(float)) / (1024 * 1024) << " MB input, " + << (n * (long long)n * sizeof(float)) / (1024 * 1024) + << " MB output)\n" << std::flush; + + cv::Mat S; + bool have_cpu = false, have_gpu = false; + cv::Mat S_cpu, S_gpu_host; + + if (mode == "cpu" || mode == "bench") { + std::cerr << "[calibration] running cv::gemm on CPU...\n" << std::flush; + auto t0 = std::chrono::steady_clock::now(); + cv::gemm(E, E, 1.0, cv::noArray(), 0.0, S_cpu, cv::GEMM_2_T); + auto t1 = std::chrono::steady_clock::now(); + double secs = std::chrono::duration<double>(t1 - t0).count(); + std::cerr << "[calibration] cv::gemm CPU: " << secs << "s for " + << n << "x" << n << " similarity matrix\n"; + have_cpu = true; + } + + if ((mode == "gpu" || mode == "bench") && have_ocl) { + std::cerr << "[calibration] running cv::gemm on OpenCL device...\n" << std::flush; + cv::UMat E_gpu, S_gpu; + E.copyTo(E_gpu); + auto t0 = std::chrono::steady_clock::now(); + cv::gemm(E_gpu, E_gpu, 1.0, cv::noArray(), 0.0, S_gpu, cv::GEMM_2_T); + S_gpu_host = S_gpu.getMat(cv::ACCESS_READ).clone(); + auto t1 = std::chrono::steady_clock::now(); + double secs = std::chrono::duration<double>(t1 - t0).count(); + std::cerr << "[calibration] cv::gemm OpenCL: " << secs << "s for " + << n << "x" << n << " similarity matrix"; + if (have_cpu) std::cerr << " max|diff|=" << cv::norm(S_cpu, S_gpu_host, cv::NORM_INF); + std::cerr << "\n"; + have_gpu = true; + } else if (mode == "gpu") { + std::cerr << "[calibration] SAE_CALIB_GEMM=gpu requested but OpenCL is unavailable\n"; + } + + if (mode == "gpu" && have_gpu) S = S_gpu_host; + else if (have_cpu) S = S_cpu; + else if (have_gpu) S = S_gpu_host; + else throw std::runtime_error("calibrate_gallery: no gemm backend produced a result"); + + // ── Histogram of pairwise similarities ────────────────────────────────── + // Instead of storing one (similarity, label) sample per pair (~n^2/2 + // entries — too many for the gradient descent below), bucket pairs into + // kHistBins bins over sim ∈ [-1, 1] and fit the sigmoid against the + // per-bin (positive_count, negative_count). + std::vector<double> pos_hist(kHistBins, 0.0), neg_hist(kHistBins, 0.0); + constexpr float bin_width = 2.f / kHistBins; + + for (int i = 0; i < n; ++i) { + const float* row = S.ptr<float>(i); + for (int j = i + 1; j < n; ++j) { + int bin = static_cast<int>((row[j] + 1.f) / bin_width); + bin = std::clamp(bin, 0, kHistBins - 1); + if (flat_actor_dedup[i] == flat_actor_dedup[j]) { + if (actor_eligible[flat_actor_dedup[i]]) pos_hist[bin] += 1.0; + // same actor but ineligible (< kMinEmbeddingsForPositive) — skip + } else { + neg_hist[bin] += 1.0; + } + } + } + + double n_pos = 0.0, n_neg = 0.0; + for (int b = 0; b < kHistBins; ++b) { n_pos += pos_hist[b]; n_neg += neg_hist[b]; } if (n_pos < 2 || n_neg < 1) { std::cerr << "[calibration] insufficient pairs (+" << n_pos @@ -69,50 +206,192 @@ inline GalleryCalibration calibrate_gallery( } // Class weights to handle pos/neg imbalance - float total = static_cast<float>(Y.size()); - float w_pos = total / (2.f * n_pos); - float w_neg = total / (2.f * n_neg); + double total = n_pos + n_neg; + double w_pos = total / (2.0 * n_pos); + double w_neg = total / (2.0 * n_neg); - // Gradient descent logistic regression (2 parameters: a, b) - float a = 10.f, b = -5.f; + std::vector<float> bin_center(kHistBins); + for (int b = 0; b < kHistBins; ++b) bin_center[b] = -1.f + (b + 0.5f) * bin_width; + + // Gradient descent logistic regression (2 parameters: a, bias) + float a = 10.f, bias = -5.f; constexpr float lr = 0.05f; constexpr int max_iter = 20000; constexpr float tol = 1e-7f; for (int iter = 0; iter < max_iter; ++iter) { - float da = 0.f, db = 0.f; - for (int i = 0; i < static_cast<int>(X.size()); ++i) { - float z = a * X[i] + b; + double da = 0.0, db = 0.0; + for (int b = 0; b < kHistBins; ++b) { + float x = bin_center[b]; + float z = a * x + bias; float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z)) : std::exp(z) / (1.f + std::exp(z)); - float err = sig - Y[i]; - float w = (Y[i] > 0.5f) ? w_pos : w_neg; - da += w * err * X[i]; - db += w * err; + if (pos_hist[b] > 0.0) { + double err = (sig - 1.f) * w_pos * pos_hist[b]; + da += err * x; + db += err; + } + if (neg_hist[b] > 0.0) { + double err = sig * w_neg * neg_hist[b]; + da += err * x; + db += err; + } } da /= total; db /= total; - a -= lr * da; - b -= lr * db; + a -= lr * static_cast<float>(da); + bias -= lr * static_cast<float>(db); if (da * da + db * db < tol * tol) break; } // Training accuracy at P=0.5 decision boundary - int correct = 0; - for (int i = 0; i < static_cast<int>(X.size()); ++i) { - float z = a * X[i] + b; + double correct = 0.0; + for (int b = 0; b < kHistBins; ++b) { + float z = a * bin_center[b] + bias; float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z)) : std::exp(z) / (1.f + std::exp(z)); - if ((sig > 0.5f) == (Y[i] > 0.5f)) ++correct; + correct += (sig > 0.5f) ? pos_hist[b] : neg_hist[b]; } - float acc = 100.f * correct / static_cast<float>(Y.size()); + double acc = 100.0 * correct / total; - GalleryCalibration cal{a, b, true}; + GalleryCalibration cal{a, bias, true}; std::cerr << "[calibration] sigmoid fitted:" - << " a=" << a << " b=" << b + << " a=" << a << " b=" << bias << " boundary(P=0.5)=sim" << cal.boundary_at(0.5f) - << " pairs=" << Y.size() - << " (+" << n_pos << "/-" << n_neg << ")" + << " pairs=" << static_cast<long long>(total) + << " (+" << static_cast<long long>(n_pos) << "/-" << static_cast<long long>(n_neg) << ")" + << " bins=" << kHistBins << " train_acc=" << acc << "%\n"; return cal; } + +// Dumps the fitted P(match | similarity) sigmoid as both a CSV table +// (similarity, p_match) and a PNG plot, to <base_path>.csv / <base_path>.png. +inline void save_calibration_curve(const GalleryCalibration& cal, const std::string& base_path) { + if (!cal.valid) return; + + constexpr int kSamples = 200; + + std::ofstream csv(base_path + ".csv"); + if (csv.is_open()) { + csv << "similarity,p_match\n"; + for (int i = 0; i <= kSamples; ++i) { + float sim = -1.f + i * (2.f / kSamples); + csv << sim << "," << cal.probability(sim) << "\n"; + } + } + + constexpr int W = 800, H = 600, M = 50; + cv::Mat img(H, W, CV_8UC3, cv::Scalar(255, 255, 255)); + cv::line(img, {M, H - M}, {W - M, H - M}, {0, 0, 0}, 1); // x-axis: similarity [-1,1] + cv::line(img, {M, M}, {M, H - M}, {0, 0, 0}, 1); // y-axis: P(match) [0,1] + + auto to_point = [&](float sim, float p) { + int x = M + static_cast<int>((sim + 1.f) * 0.5f * (W - 2 * M)); + int y = (H - M) - static_cast<int>(p * (H - 2 * M)); + return cv::Point(x, y); + }; + + // P=0.5 reference line and decision boundary + float boundary = cal.boundary_at(0.5f); + cv::line(img, to_point(-1.f, 0.5f), to_point(1.f, 0.5f), {200, 200, 200}, 1); + if (boundary >= -1.f && boundary <= 1.f) + cv::line(img, to_point(boundary, 0.f), to_point(boundary, 1.f), {200, 200, 200}, 1); + + cv::Point prev = to_point(-1.f, cal.probability(-1.f)); + for (int i = 1; i <= kSamples; ++i) { + float sim = -1.f + i * (2.f / kSamples); + cv::Point pt = to_point(sim, cal.probability(sim)); + cv::line(img, prev, pt, {255, 0, 0}, 2); + prev = pt; + } + + cv::imwrite(base_path + ".png", img); + std::cerr << "[calibration] saved curve to " << base_path << ".csv / " << base_path << ".png\n"; +} + +// FNV-1a 64-bit hash over the gallery's reference embeddings and actor +// assignments, used to detect whether a cached calibration is still valid. +inline uint64_t hash_gallery_embeddings( + const std::vector<Embedding>& flat_emb, + const std::vector<int>& flat_actor) +{ + uint64_t h = 1469598103934665603ULL; + constexpr uint64_t prime = 1099511628211ULL; + auto mix = [&](const void* data, size_t n) { + const auto* p = static_cast<const unsigned char*>(data); + for (size_t i = 0; i < n; ++i) { + h ^= p[i]; + h *= prime; + } + }; + + uint64_t n = flat_emb.size(); + mix(&n, sizeof(n)); + for (const auto& emb : flat_emb) + mix(emb.data(), emb.size() * sizeof(float)); + for (int a : flat_actor) + mix(&a, sizeof(a)); + return h; +} + +// Calibrates the gallery, caching the fitted (a, b, valid) result on disk +// keyed by a hash of the reference embeddings. The O(n^2) pairwise fit only +// re-runs when the gallery's embeddings/actor assignments actually change. +inline GalleryCalibration calibrate_gallery_cached( + const std::vector<Embedding>& flat_emb, + const std::vector<int>& flat_actor, + const std::string& cache_path) +{ + uint64_t hash = hash_gallery_embeddings(flat_emb, flat_actor); + + std::string base_path = cache_path; + constexpr std::string_view kJsonExt = ".json"; + if (base_path.size() >= kJsonExt.size() && + base_path.compare(base_path.size() - kJsonExt.size(), kJsonExt.size(), kJsonExt) == 0) + base_path.resize(base_path.size() - kJsonExt.size()); + + std::ifstream in(cache_path); + if (in.is_open()) { + try { + nlohmann::json j; + in >> j; + if (j.at("hash").get<uint64_t>() == hash) { + GalleryCalibration cal; + cal.a = j.at("a").get<float>(); + cal.b = j.at("b").get<float>(); + cal.valid = j.at("valid").get<bool>(); + std::cerr << "[calibration] using cached calibration from " + << cache_path << " (a=" << cal.a << " b=" << cal.b + << " valid=" << cal.valid << ")\n"; + save_calibration_curve(cal, base_path); + return cal; + } + std::cerr << "[calibration] cache at " << cache_path + << " is stale, recomputing\n"; + } catch (const std::exception&) { + std::cerr << "[calibration] cache at " << cache_path + << " is unreadable, recomputing\n"; + } + } + + auto t0 = std::chrono::steady_clock::now(); + GalleryCalibration cal = calibrate_gallery(flat_emb, flat_actor); + auto t1 = std::chrono::steady_clock::now(); + double secs = std::chrono::duration<double>(t1 - t0).count(); + std::cerr << "[calibration] fit took " << secs << "s for " + << flat_emb.size() << " embeddings\n"; + + nlohmann::json j; + j["hash"] = hash; + j["a"] = cal.a; + j["b"] = cal.b; + j["valid"] = cal.valid; + j["fit_secs"] = secs; + std::ofstream out(cache_path); + if (out.is_open()) out << j.dump(2) << "\n"; + + save_calibration_curve(cal, base_path); + + return cal; +} diff --git a/src/gallery/gallery_store.cpp b/src/gallery/gallery_store.cpp index 053321a..1e688a1 100644 --- a/src/gallery/gallery_store.cpp +++ b/src/gallery/gallery_store.cpp @@ -1,7 +1,9 @@ #include "gallery_store.hpp" #include <nlohmann/json.hpp> +#include <chrono> #include <fstream> +#include <iostream> #include <stdexcept> using json = nlohmann::json; @@ -11,14 +13,24 @@ ActorGallery load_gallery(const std::string& path) { if (!f.is_open()) throw std::runtime_error("load_gallery: cannot open " + path); + std::cerr << "[gallery] loading " << path << "..." << std::flush; + auto t0 = std::chrono::steady_clock::now(); + json j; f >> j; + auto t1 = std::chrono::steady_clock::now(); + std::cerr << " parsed JSON in " + << std::chrono::duration<double>(t1 - t0).count() << "s\n"; + ActorGallery gallery; for (const auto& ja : j.at("actors")) { ActorGallery::Actor actor; - actor.imdb_id = ja.at("imdb_id").get<std::string>(); - actor.name = ja.at("name").get<std::string>(); + actor.imdb_id = ja.value("imdb_id", ""); + actor.tmdb_id = ja.value("tmdb_id", ""); + // older make_jellyfin_gallery.py galleries used "jellyfin_person_id" + actor.jellyfin_id = ja.value("jellyfin_id", ja.value("jellyfin_person_id", "")); + actor.name = ja.at("name").get<std::string>(); if (ja.contains("source_images")) actor.source_images = ja.at("source_images").get<std::vector<std::string>>(); @@ -30,6 +42,14 @@ ActorGallery load_gallery(const std::string& path) { gallery.actors.push_back(std::move(actor)); } + + size_t n_emb = 0; + for (const auto& actor : gallery.actors) n_emb += actor.embeddings.size(); + auto t2 = std::chrono::steady_clock::now(); + std::cerr << "[gallery] built " << gallery.actors.size() << " actors / " + << n_emb << " embeddings in " + << std::chrono::duration<double>(t2 - t1).count() << "s\n"; + return gallery; } @@ -40,6 +60,8 @@ void save_gallery(const std::string& path, const ActorGallery& gallery) { for (const auto& actor : gallery.actors) { json ja; ja["imdb_id"] = actor.imdb_id; + ja["tmdb_id"] = actor.tmdb_id; + ja["jellyfin_id"] = actor.jellyfin_id; ja["name"] = actor.name; ja["source_images"] = actor.source_images; diff --git a/src/gallery/gallery_store.hpp b/src/gallery/gallery_store.hpp index 24d0213..b7c29dc 100644 --- a/src/gallery/gallery_store.hpp +++ b/src/gallery/gallery_store.hpp @@ -8,7 +8,9 @@ // { // "actors": [ // { -// "imdb_id": "nm0000093", +// "imdb_id": "nm0000093", // optional, "" if unknown +// "tmdb_id": "287", // optional, "" if unknown +// "jellyfin_id": "abc123-guid", // optional, "" unless from make_jellyfin_gallery.py // "name": "Brad Pitt", // "source_images": ["img1.jpg", "img2.jpg"], // "embeddings": [[0.012, -0.034, ...], ...] // one 512-float array per image diff --git a/src/nodes/identity_matcher_node.hpp b/src/nodes/identity_matcher_node.hpp index 53ed2b2..ddd337f 100644 --- a/src/nodes/identity_matcher_node.hpp +++ b/src/nodes/identity_matcher_node.hpp @@ -37,6 +37,7 @@ struct IdentityMatcherFunc { , ratio_(cfg.match_ratio) , ratio_ceil_(cfg.match_ratio_ceil) { + std::cerr << "[identity_matcher] flattening gallery embeddings...\n"; for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) { for (const auto& emb : gallery_.actors[ai].embeddings) { flat_emb_.push_back(emb); @@ -44,7 +45,10 @@ struct IdentityMatcherFunc { } } - cal_ = calibrate_gallery(flat_emb_, flat_actor_); + std::cerr << "[identity_matcher] starting calibration (" + << flat_emb_.size() << " embeddings)...\n"; + cal_ = calibrate_gallery_cached(flat_emb_, flat_actor_, + cfg.gallery_path + ".calib_cache.json"); if (cal_.valid) { std::cerr << "[identity_matcher] calibrated Bayesian matching" @@ -126,9 +130,11 @@ struct IdentityMatcherFunc { ia.track_id = tf.track_ids[fi]; if (accept) { - ia.actor_idx = best_actor; - ia.name = gallery_.actors[best_actor].name; - ia.imdb_id = gallery_.actors[best_actor].imdb_id; + ia.actor_idx = best_actor; + ia.name = gallery_.actors[best_actor].name; + ia.imdb_id = gallery_.actors[best_actor].imdb_id; + ia.tmdb_id = gallery_.actors[best_actor].tmdb_id; + ia.jellyfin_id = gallery_.actors[best_actor].jellyfin_id; ia.similarity = cal_.valid ? cal_.probability(best_s, log_prior_odds_) : best_s; diff --git a/src/nodes/result_sink_node.hpp b/src/nodes/result_sink_node.hpp index 2c9c416..7afad22 100644 --- a/src/nodes/result_sink_node.hpp +++ b/src/nodes/result_sink_node.hpp @@ -18,7 +18,16 @@ using json = nlohmann::json; // KPN sink node: accumulates SceneAnnotations and writes the final JSON on EOF. // // Verbosity::minimal — merges per-frame presence into contiguous time windows. -// Output: { "movie": "...", "actors": [{ "name", "imdb_id", "scenes": [[t0,t1], ...] }] } +// Output: { +// "schema_version": 1, "movie": "...", "sample_fps": ..., "anneal_sec": ..., +// "actors": [{ "name", "imdb_id", "tmdb_id", "jellyfin_id", "scenes": [[t0,t1], ...] }] +// } +// This is the spec consumed by the Jellyfin plugin: each actor carries every +// identity key the gallery knows (empty string if not resolved). The plugin +// should prefer "jellyfin_id" (direct Person item GUID) when non-empty, and +// otherwise resolve "imdb_id"/"tmdb_id" against the item's People ProviderIds. +// To find who's on screen at timestamp t, scan each actor's "scenes" for a +// window where start <= t <= end. // // Verbosity::standard — per-frame detail including bboxes, similarity, unknowns. // Output: { "frames": [{ "t", "identified": [...], "unknowns": [...] }] } @@ -55,6 +64,10 @@ struct ResultSinkFunc { } private: + // Bump when the minimal/standard output JSON structure changes in a way + // the Jellyfin plugin needs to detect. + static constexpr int kSchemaVersion = 1; + static int count_known(const std::vector<IdentifiedActor>& v) { int n = 0; for (const auto& a : v) if (a.actor_idx >= 0) ++n; @@ -73,6 +86,7 @@ private: if (cfg_.verbosity == Verbosity::xray) { root = build_xray(); } else { + root["schema_version"] = kSchemaVersion; root["movie"] = cfg_.movie_path; root["sample_fps"] = cfg_.sample_fps; root["anneal_sec"] = cfg_.anneal_sec; @@ -91,20 +105,20 @@ private: } struct ActorWindow { - std::string name, imdb_id; + std::string name, imdb_id, tmdb_id, jellyfin_id; std::vector<std::pair<double, double>> scenes; // [start_sec, end_sec] }; // Core logic: merge per-frame detections into annealed [start, end] windows. std::vector<ActorWindow> build_actor_windows() { - struct Info { std::string name, imdb_id; }; + struct Info { std::string name, imdb_id, tmdb_id, jellyfin_id; }; std::map<int, Info> actor_info; std::map<int, std::vector<double>> timestamps; for (const auto& frame : frames_) { for (const auto& ia : frame.visible_actors) { if (ia.actor_idx < 0) continue; - actor_info[ia.actor_idx] = {ia.name, ia.imdb_id}; + actor_info[ia.actor_idx] = {ia.name, ia.imdb_id, ia.tmdb_id, ia.jellyfin_id}; timestamps[ia.actor_idx].push_back(frame.timestamp_sec); } } @@ -112,8 +126,10 @@ private: std::vector<ActorWindow> result; for (auto& [idx, ts_vec] : timestamps) { ActorWindow aw; - aw.name = actor_info[idx].name; - aw.imdb_id = actor_info[idx].imdb_id; + aw.name = actor_info[idx].name; + aw.imdb_id = actor_info[idx].imdb_id; + aw.tmdb_id = actor_info[idx].tmdb_id; + aw.jellyfin_id = actor_info[idx].jellyfin_id; double win_start = ts_vec[0], win_end = ts_vec[0]; for (size_t i = 1; i < ts_vec.size(); ++i) { @@ -136,9 +152,11 @@ private: for (const auto& [s, e] : aw.scenes) windows.push_back({s, e}); json ja; - ja["name"] = aw.name; - ja["imdb_id"] = aw.imdb_id; - ja["scenes"] = std::move(windows); + ja["name"] = aw.name; + ja["imdb_id"] = aw.imdb_id; + ja["tmdb_id"] = aw.tmdb_id; + ja["jellyfin_id"] = aw.jellyfin_id; + ja["scenes"] = std::move(windows); actors.push_back(std::move(ja)); } return actors; @@ -178,8 +196,10 @@ private: if (ia.actor_idx >= 0) { json ja; - ja["name"] = ia.name; - ja["imdb_id"] = ia.imdb_id; + ja["name"] = ia.name; + ja["imdb_id"] = ia.imdb_id; + ja["tmdb_id"] = ia.tmdb_id; + ja["jellyfin_id"] = ia.jellyfin_id; ja["similarity"] = ia.similarity; ja["track_id"] = ia.track_id; ja["bbox"] = jbox; diff --git a/src/nodes/scene_tracker_node.hpp b/src/nodes/scene_tracker_node.hpp index 73d4198..5ca6e31 100644 --- a/src/nodes/scene_tracker_node.hpp +++ b/src/nodes/scene_tracker_node.hpp @@ -41,6 +41,8 @@ struct SceneTrackerFunc { slot.last_crop = ia.crop; slot.name = ia.name; slot.imdb_id = ia.imdb_id; + slot.tmdb_id = ia.tmdb_id; + slot.jellyfin_id = ia.jellyfin_id; // Keep the best (highest) similarity seen in this window if (ia.similarity > slot.best_similarity) slot.best_similarity = ia.similarity; @@ -60,9 +62,11 @@ struct SceneTrackerFunc { for (const auto& [actor_idx, slot] : active_) { IdentifiedActor ia; - ia.actor_idx = actor_idx; - ia.name = slot.name; - ia.imdb_id = slot.imdb_id; + ia.actor_idx = actor_idx; + ia.name = slot.name; + ia.imdb_id = slot.imdb_id; + ia.tmdb_id = slot.tmdb_id; + ia.jellyfin_id = slot.jellyfin_id; ia.similarity = slot.best_similarity; ia.bbox = slot.last_bbox; ia.crop = slot.last_crop; @@ -85,6 +89,8 @@ private: cv::Mat last_crop; std::string name; std::string imdb_id; + std::string tmdb_id; + std::string jellyfin_id; }; double extinction_sec_; diff --git a/src/python_bindings.cpp b/src/python_bindings.cpp new file mode 100644 index 0000000..a9623e5 --- /dev/null +++ b/src/python_bindings.cpp @@ -0,0 +1,40 @@ +// sae_embed — Python module wrapping FaceEmbedderEngine (SCRFD + ArcFace). +// +// Loads both ONNX sessions once per FaceEmbedder instance, then embeds many +// images via repeated embed() calls — avoiding the per-process model-load +// cost of the embed_faces CLI when embedding a large gallery. + +#include "face_embedder_engine.hpp" + +#include <nanobind/nanobind.h> +#include <nanobind/stl/optional.h> +#include <nanobind/stl/string.h> +#include <nanobind/stl/vector.h> + +namespace nb = nanobind; +using namespace nb::literals; + +NB_MODULE(sae_embed, m) { + m.doc() = "SCRFD + ArcFace face embedding, models loaded once per FaceEmbedder"; + + nb::class_<FaceEmbedResult>(m, "FaceResult") + .def_ro("ok", &FaceEmbedResult::ok) + .def_ro("error", &FaceEmbedResult::error) + .def_ro("confidence", &FaceEmbedResult::confidence) + .def_prop_ro("embedding", [](const FaceEmbedResult& r) -> std::optional<std::vector<float>> { + if (!r.ok) return std::nullopt; + return std::vector<float>(r.embedding.begin(), r.embedding.end()); + }) + .def_prop_ro("bbox", [](const FaceEmbedResult& r) { + return std::vector<float>{r.bbox[0], r.bbox[1], r.bbox[2], r.bbox[3]}; + }); + + nb::class_<FaceEmbedderEngine>(m, "FaceEmbedder") + .def(nb::init<std::string, std::string, float, float, int>(), + "detector_model"_a, "arcface_model"_a, + "conf"_a = 0.5f, "nms"_a = 0.4f, "max_side"_a = 500) + .def("embed", &FaceEmbedderEngine::embed_path, "path"_a, + nb::call_guard<nb::gil_scoped_release>(), + "Detect the highest-confidence face in the image, align it, and " + "return a FaceResult with its 512-d ArcFace embedding."); +} diff --git a/src/types.hpp b/src/types.hpp index 35f2b67..73c26e5 100644 --- a/src/types.hpp +++ b/src/types.hpp @@ -92,6 +92,8 @@ struct IdentifiedActor { int track_id{-1}; // face track ID from FaceTrackerFunc std::string name; std::string imdb_id; + std::string tmdb_id; + std::string jellyfin_id; // Jellyfin Person item GUID, if gallery was built from Jellyfin float similarity{0.f}; // calibrated P(match) or cosine similarity; 0 for unknowns cv::Rect2f bbox; cv::Mat crop; // 112×112 aligned crop (stored as shared_ptr by KPN) @@ -116,6 +118,8 @@ struct SceneAnnotation { struct ActorGallery { struct Actor { std::string imdb_id; + std::string tmdb_id; + std::string jellyfin_id; // Jellyfin Person item GUID, if known std::string name; std::vector<Embedding> embeddings; // one per reference image std::vector<std::string> source_images;