diff --git a/README.md b/README.md index c8e7312..b05dfed 100644 --- a/README.md +++ b/README.md @@ -151,6 +151,37 @@ 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. +After a successful run, the output JSON is pushed to the JRay Jellyfin +plugin's Truth endpoint (`PUT /Plugins/JRay/Items/{itemId}/Truth`) so +Jellyfin picks it up immediately, using `--api-key` (must be an +**Administrator** key for the push to succeed). Pass `--no-push` to skip +this and only write `--output` locally (e.g. for local debugging). + +### Worker mode + +Pass `--worker` instead of `--item-id`/`--title` to run this as an extraction +worker: it polls the JRay plugin's `GET /Plugins/JRay/Tasks/Pending` endpoint +for a random batch of items with no truth data yet, processes each one, and +pushes the result back. The endpoint's sampling spreads work across the +backlog without any server-side task tracking, so any number of workers can +poll the same library concurrently. + +```bash +python3 scripts/run_from_jellyfin.py \ + --jellyfin-url http://jellyfin.local:8096 \ + --api-key \ + --gallery whole_gallery.json \ + --worker \ + -- --fps 5 +``` + +- `--poll-limit` — batch size requested from `Tasks/Pending` (default 10, max 100) +- `--poll-interval` — seconds to sleep between polls when the backlog is empty (default 60) +- `--once` — process a single batch and exit instead of looping forever + +A failure on one item (bad path, push rejected, etc.) is logged and the +worker moves on to the next item rather than exiting. + ## Output format **Minimal** (default) — Jellyfin-ready: diff --git a/scripts/filter_gallery.py b/scripts/filter_gallery.py index a516649..7499c97 100644 --- a/scripts/filter_gallery.py +++ b/scripts/filter_gallery.py @@ -64,8 +64,26 @@ def find_item_id(base_url: str, api_key: str, title: str, item_types: list[str]) 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"} + # /Items/{id} (no user context) returns 400 on recent Jellyfin servers; + # /Items?Ids=... works with an API key alone. + data = jf_get(base_url, api_key, "/Items", Ids=item_id, Fields="People", Recursive="true") + items = data.get("Items", []) + if not items: + raise ValueError(f"Item {item_id} not found") + item = items[0] + cast_ids = {p["Id"] for p in item.get("People", []) if p.get("Type") == "Actor"} + + # Episodes generally only carry their own guest stars in "People" — the + # regular/recurring cast lives on the parent Series item, so pull that + # in too or the filtered gallery ends up missing the main cast. + series_id = item.get("SeriesId") + if series_id: + series_data = jf_get(base_url, api_key, "/Items", Ids=series_id, Fields="People", Recursive="true") + series_items = series_data.get("Items", []) + if series_items: + cast_ids |= {p["Id"] for p in series_items[0].get("People", []) if p.get("Type") == "Actor"} + + return cast_ids def main(): diff --git a/scripts/make_gallery.py b/scripts/make_gallery.py index fe04554..aeb762d 100755 --- a/scripts/make_gallery.py +++ b/scripts/make_gallery.py @@ -24,7 +24,6 @@ Usage: # Additional options: # --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 # --image-dir /tmp/gallery_imgs where to cache downloaded images @@ -46,6 +45,11 @@ from sae_embed_loader import load_embedder TMDB_BASE = "https://api.themoviedb.org/3" TMDB_IMG = "https://image.tmdb.org/t/p/original" +WIKIDATA_SPARQL = "https://query.wikidata.org/sparql" +WIKIDATA_HEADERS = { + "Accept": "application/sparql-results+json", + "User-Agent": "scene-actor-extraction/1.0 (https://github.com/; gallery builder)", +} # ── TMDB helpers ────────────────────────────────────────────────────────────── @@ -61,6 +65,29 @@ def tmdb_get(path: str, token: str, **params) -> dict: return r.json() +# ── Wikidata helpers ────────────────────────────────────────────────────────── + +def wikidata_image_urls(imdb_person_id: str) -> list[str]: + """Return Commons image URL(s) for a person via their IMDB ID (P345 -> P18).""" + if not imdb_person_id: + return [] + query = ( + "SELECT ?image WHERE { " + f'?person wdt:P345 "{imdb_person_id}" . ' + "?person wdt:P18 ?image . " + "}" + ) + try: + r = requests.get(WIKIDATA_SPARQL, params={"query": query, "format": "json"}, + headers=WIKIDATA_HEADERS, timeout=15) + r.raise_for_status() + bindings = r.json().get("results", {}).get("bindings", []) + return [b["image"]["value"] for b in bindings if "image" in b] + except Exception as e: + print(f" [warn] wikidata lookup failed for {imdb_person_id}: {e}", file=sys.stderr) + return [] + + def tmdb_id_from_imdb(imdb_id: str, key: str) -> int: data = tmdb_get(f"/find/{imdb_id}", key, external_source="imdb_id") results = data.get("movie_results", []) @@ -69,10 +96,10 @@ def tmdb_id_from_imdb(imdb_id: str, key: str) -> int: return results[0]["id"] -def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]: +def fetch_cast(movie_id: int, key: str) -> list[dict]: """Return list of {id, name, imdb_id, profile_images: [...url...]}.""" credits = tmdb_get(f"/movie/{movie_id}/credits", key) - cast = credits.get("cast", [])[:max_actors] + cast = credits.get("cast", []) actors = [] for member in cast: @@ -88,8 +115,13 @@ def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]: image_urls = [TMDB_IMG + p["file_path"] for p in profiles if p.get("file_path")] if not image_urls: - print(f" [warn] no images for {member['name']}, skipping", file=sys.stderr) - continue + image_urls = wikidata_image_urls(imdb_id) + if image_urls: + print(f" [info] no TMDB images for {member['name']}, " + f"found {len(image_urls)} via Wikidata", file=sys.stderr) + + if not image_urls: + print(f" [warn] no images for {member['name']}", file=sys.stderr) actors.append({ "id": person_id, @@ -127,14 +159,15 @@ def download_images(actor: dict, dest_dir: Path, n: int) -> list[Path]: # ── Gallery assembly ───────────────────────────────────────────────────────── 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.""" + images_per_actor: int, + image_root: Path) -> tuple[dict, list[dict]]: + """Fetch cast, download images, embed, return (gallery dict, actors needing more images).""" print(f"Fetching cast for TMDB movie {movie_id}…", file=sys.stderr) - actors = fetch_cast(movie_id, key, max_actors) - print(f"Found {len(actors)} actors with images", file=sys.stderr) + actors = fetch_cast(movie_id, key) + print(f"Found {len(actors)} cast member(s)", file=sys.stderr) gallery_actors = [] + missing = [] for actor in actors: safe_name = actor["name"].replace(" ", "_") @@ -142,9 +175,18 @@ def build_gallery(movie_id: int, key: str, embedder, actor_dir = image_root / f"{dir_id}_{safe_name}" print(f"\n{actor['name']} ({dir_id})", file=sys.stderr) + + if not actor["profile_images"]: + print(" no images found, skipping", file=sys.stderr) + missing.append({"name": actor["name"], "imdb_id": actor["imdb_id"], + "tmdb_id": actor["tmdb_id"], "reason": "no images found"}) + continue + image_paths = download_images(actor, actor_dir, images_per_actor) if not image_paths: print(" no images downloaded, skipping", file=sys.stderr) + missing.append({"name": actor["name"], "imdb_id": actor["imdb_id"], + "tmdb_id": actor["tmdb_id"], "reason": "download failed"}) continue print(f" embedding {len(image_paths)} image(s)…", file=sys.stderr) @@ -163,6 +205,8 @@ def build_gallery(movie_id: int, key: str, embedder, if not embeddings: print(" no valid embeddings, skipping actor", file=sys.stderr) + missing.append({"name": actor["name"], "imdb_id": actor["imdb_id"], + "tmdb_id": actor["tmdb_id"], "reason": "no valid embeddings"}) continue gallery_actors.append({ @@ -175,7 +219,7 @@ def build_gallery(movie_id: int, key: str, embedder, }) print(f" → {len(embeddings)} embedding(s) stored", file=sys.stderr) - return {"actors": gallery_actors} + return {"actors": gallery_actors}, missing # ── Entry point ─────────────────────────────────────────────────────────────── @@ -197,8 +241,6 @@ def main(): 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("--max-actors", type=int, default=20, - help="Maximum number of cast members to include (default: 20)") parser.add_argument("--images-per-actor",type=int, default=3, help="Profile images to download per actor (default: 3)") parser.add_argument("--image-dir", default=None, @@ -221,11 +263,10 @@ def main(): print(f"TMDB movie ID: {movie_id}", file=sys.stderr) # Build gallery - gallery = build_gallery( + gallery, missing = build_gallery( movie_id = movie_id, key = args.tmdb_key, embedder = embedder, - max_actors = args.max_actors, images_per_actor = args.images_per_actor, image_root = image_root, ) @@ -242,6 +283,11 @@ def main(): output.write_text(json.dumps(gallery, indent=2) + "\n") print(f"Saved: {output}", file=sys.stderr) + if missing: + missing_path = output.with_name(output.stem + ".missing_images.json") + missing_path.write_text(json.dumps(missing, indent=2) + "\n") + print(f"{len(missing)} actor(s) need images — see {missing_path}", file=sys.stderr) + if __name__ == "__main__": main() diff --git a/scripts/make_jellyfin_gallery.py b/scripts/make_jellyfin_gallery.py index ac6c8a7..b382490 100644 --- a/scripts/make_jellyfin_gallery.py +++ b/scripts/make_jellyfin_gallery.py @@ -261,7 +261,7 @@ def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict, 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: + embed_executor: concurrent.futures.ThreadPoolExecutor) -> tuple[dict | None, str | 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 @@ -274,7 +274,7 @@ def embed_actor(pid: str, info: dict, image_paths: list[Path], name = info["name"] if not image_paths: print(f" {name}: no image available, skipping", file=sys.stderr) - return None + return None, "no image available" embeddings = [] source_images = [] @@ -289,7 +289,7 @@ def embed_actor(pid: str, info: dict, image_paths: list[Path], if not embeddings: print(f" {name}: no valid embeddings, skipping actor", file=sys.stderr) - return None + return None, "no valid embeddings" print(f" {name}: → {len(embeddings)} embedding(s) stored", file=sys.stderr) return { @@ -300,13 +300,13 @@ def embed_actor(pid: str, info: dict, image_paths: list[Path], "source_images": source_images, "embeddings": embeddings, "appearances": info["appearances"], - } + }, None 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: + embed_executor: concurrent.futures.ThreadPoolExecutor) -> tuple[dict | None, str | None]: safe_name = info["name"].replace(" ", "_") actor_dir = image_root / f"{pid}_{safe_name}" image_paths, imdb_id, tmdb_id = fetch_actor_images( @@ -319,44 +319,54 @@ def process_actor(pid: str, info: dict, base_url: str, api_key: str, 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: + tmdb_key: str | None = None, workers: int = 8) -> tuple[dict, list[dict]]: actors = collect_actors(base_url, api_key, item_types) gallery_actors = [] todo = [] + n_retry = 0 for pid, info in actors.items(): - if pid in existing_actors: - gallery_actors.append(existing_actors[pid]) + existing = existing_actors.get(pid) + if existing is not None and existing.get("embeddings"): + gallery_actors.append(existing) else: + if existing is not None: + n_retry += 1 todo.append((pid, info)) if len(gallery_actors): print(f"Skipping {len(gallery_actors)} actor(s) already present in existing gallery", file=sys.stderr) + if n_retry: + print(f"Retrying {n_retry} actor(s) with no embeddings in existing gallery", file=sys.stderr) print(f"Processing {len(todo)} new actor(s) with {workers} worker(s)…", file=sys.stderr) + missing = [] 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"] + fetch_imdb_ids, tmdb_key, embed_executor): (pid, info["name"]) for pid, info in todo } for future in concurrent.futures.as_completed(futures): n_done += 1 - name = futures[future] + pid, name = futures[future] try: - actor = future.result() + actor, reason = future.result() except Exception as e: print(f" [error] {name}: {e}", file=sys.stderr) + missing.append({"jellyfin_id": pid, "name": name, "reason": str(e)}) continue if actor: gallery_actors.append(actor) + else: + missing.append({"jellyfin_id": pid, "name": name, "reason": reason}) if n_done % 25 == 0: print(f" progress: {n_done}/{len(todo)} actors processed", file=sys.stderr) - return {"actors": gallery_actors} + return {"actors": gallery_actors}, missing # ── Entry point ─────────────────────────────────────────────────────────────── @@ -423,7 +433,7 @@ def main(): existing_actors[pid] = actor print(f"Loaded {len(existing_actors)} actor(s) from existing gallery for merge", file=sys.stderr) - gallery = build_gallery( + gallery, missing = build_gallery( base_url=jellyfin_url, api_key=args.api_key, embedder=embedder, @@ -447,6 +457,11 @@ def main(): output.write_text(json.dumps(gallery, indent=2) + "\n") print(f"Saved: {output}", file=sys.stderr) + if missing: + missing_path = output.with_name(output.stem + ".missing_images.json") + missing_path.write_text(json.dumps(missing, indent=2) + "\n") + print(f"{len(missing)} actor(s) need images — see {missing_path}", file=sys.stderr) + if __name__ == "__main__": main() diff --git a/scripts/run_from_jellyfin.py b/scripts/run_from_jellyfin.py index b1e7bfc..c06cade 100644 --- a/scripts/run_from_jellyfin.py +++ b/scripts/run_from_jellyfin.py @@ -15,6 +15,16 @@ Usage: -- --fps 5 --verbosity 2 Anything after "--" is passed through unchanged to scene_analyze. + +Worker mode (--worker) polls the JRay plugin's Tasks/Pending endpoint for a +random batch of items with no truth data yet, processing each in turn: + + python scripts/run_from_jellyfin.py \\ + --jellyfin-url http://jellyfin.local:8096 \\ + --api-key YOUR_API_KEY \\ + --gallery whole_gallery.json \\ + --worker \\ + -- --fps 5 """ import argparse @@ -22,57 +32,71 @@ import json import subprocess import sys import tempfile +import time from pathlib import Path +import requests + 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") + data = jf_get(base_url, api_key, "/Items", Ids=item_id, Fields="Path", Recursive="true") + items = data.get("Items", []) + if not items: + raise ValueError(f"Item {item_id} not found") + path = items[0].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 + return items[0].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:] +def fetch_pending_tasks(base_url: str, api_key: str, limit: int = 10) -> list[dict]: + """GET /Plugins/JRay/Tasks/Pending — a random batch of items with no truth data yet. - 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) + Returns a list of {"item_id", "path", "name"}; an empty list means + there's nothing left to do (for now). + """ + return jf_get(base_url, api_key, "/Plugins/JRay/Tasks/Pending", limit=limit) + +def push_truth(base_url: str, api_key: str, item_id: str, output_path: str) -> None: + """PUT the scene_analyze output JSON to the JRay plugin's Truth endpoint. + + Requires an Administrator API key (reuses --api-key). + """ + url = base_url.rstrip("/") + f"/Plugins/JRay/Items/{item_id}/Truth" + headers = {"X-Emby-Token": api_key, "Content-Type": "application/json"} + + # --verbosity 1 (standard) adds a "frames" array with per-frame bbox/ + # similarity data for local debugging — the Truth schema only needs + # schema_version/movie/sample_fps/anneal_sec/actors, and "frames" can + # be tens of MB for a full episode, well past typical proxy body limits. + payload = json.loads(Path(output_path).read_text()) + if payload.pop("frames", None) is not None: + print(f"Stripped per-frame data before push ({output_path} keeps it locally)", file=sys.stderr) + body = json.dumps(payload) + + r = requests.put(url, headers=headers, data=body, timeout=30) + if r.status_code == 204: + print(f"Pushed {output_path} -> {url}", file=sys.stderr) + return + if r.status_code in (401, 403): + raise RuntimeError(f"Push failed ({r.status_code}): API key needs Administrator rights for {url}") + if r.status_code == 400: + raise RuntimeError(f"Push failed (400 Bad Request): {r.text}") + r.raise_for_status() + + +def process_item(args, item_id: str, extra: list[str]) -> None: + """Resolve, analyze, and (unless --no-push) push truth for one Jellyfin item.""" 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)") + raise RuntimeError(f"Resolved path does not exist on this filesystem: {movie_path} " + f"(this tool must share Jellyfin's media mount)") output = args.output or f"{name}.json" @@ -91,8 +115,11 @@ def main(): filtered_file.close() gallery_path = filtered_file.name + # Default to minimal output (just actor/scene windows, no per-frame + # data) — that's all the Truth push needs. Placed before *extra so a + # debug run can override with `-- --verbosity 1` (last flag wins). cmd = [args.bin, "--movie", movie_path, "--gallery", gallery_path, - "--output", output, *extra] + "--output", output, "--verbosity", "0", *extra] print("Running:", " ".join(cmd), file=sys.stderr) if args.dry_run: @@ -104,6 +131,104 @@ def main(): if filtered_file is not None: Path(filtered_file.name).unlink(missing_ok=True) + if not args.no_push: + push_truth(args.jellyfin_url, args.api_key, item_id, output) + + +def run_worker(args, extra: list[str]) -> None: + """Poll /Plugins/JRay/Tasks/Pending and process items until the backlog is empty. + + The endpoint returns a random sample, so repeated polling naturally + spreads work across the backlog without server-side task tracking. An + empty array means there's nothing left to do (for now). + """ + while True: + try: + tasks = fetch_pending_tasks(args.jellyfin_url, args.api_key, args.poll_limit) + except requests.HTTPError as e: + print(f"Failed to fetch pending tasks: {e}", file=sys.stderr) + tasks = [] + + if not tasks: + print("No pending tasks.", file=sys.stderr) + if args.once: + return + time.sleep(args.poll_interval) + continue + + for task in tasks: + item_id = task["item_id"] + label = task.get("name") or task.get("path") or item_id + print(f"=== {label} ({item_id}) ===", file=sys.stderr) + try: + process_item(args, item_id, extra) + except Exception as e: + print(f"Failed to process {label!r} ({item_id}): {e}", file=sys.stderr) + + if args.once: + return + + +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() + group.add_argument("--item-id", help="Jellyfin item id of the title") + group.add_argument("--title", help="Title to search for (uses first match)") + group.add_argument("--worker", action="store_true", + help="Poll /Plugins/JRay/Tasks/Pending for work instead of " + "processing a single --item-id/--title") + parser.add_argument("--poll-limit", type=int, default=10, + help="Tasks/Pending batch size for --worker (default: 10, max 100)") + parser.add_argument("--poll-interval", type=float, default=60, + help="Seconds to wait between polls when --worker finds nothing " + "(default: 60)") + parser.add_argument("--once", action="store_true", + help="With --worker, process one batch and exit instead of " + "looping forever") + 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: <title>.json; " + "ignored with --worker, which always uses <title>.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") + parser.add_argument("--no-push", action="store_true", + help="Don't push results to the JRay plugin after the run " + "(local-only/debug; --output is still written)") + args, extra = parser.parse_known_args() + if extra and extra[0] == "--": + extra = extra[1:] + + if args.worker: + if args.output: + sys.exit("--output is incompatible with --worker (each item needs its own file)") + run_worker(args, extra) + return + + if not args.item_id and not args.title: + sys.exit("one of --item-id, --title, or --worker is required") + + 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) + + try: + process_item(args, item_id, extra) + except RuntimeError as e: + sys.exit(str(e)) + if __name__ == "__main__": main() diff --git a/src/embed_faces.cpp b/src/embed_faces.cpp index ff45b00..1a1241e 100644 --- a/src/embed_faces.cpp +++ b/src/embed_faces.cpp @@ -129,6 +129,13 @@ static FaceResult process(const std::string& path, std::vector<DetectedFace> faces = detect(img); + if (faces.empty()) { + cv::Mat enhanced = enhance_for_retry(img); + faces = detect(enhanced); + if (!faces.empty()) + img = enhanced; + } + if (faces.empty()) { res.error = "no face detected"; return res; diff --git a/src/face_embedder_engine.hpp b/src/face_embedder_engine.hpp index ab831ae..e89e214 100644 --- a/src/face_embedder_engine.hpp +++ b/src/face_embedder_engine.hpp @@ -66,6 +66,12 @@ public: } std::vector<DetectedFace> faces = detector_->detect(img); + if (faces.empty()) { + cv::Mat enhanced = enhance_for_retry(img); + faces = detector_->detect(enhanced); + if (!faces.empty()) + img = enhanced; + } if (faces.empty()) { res.error = "no face detected"; return res; diff --git a/src/face_utils.hpp b/src/face_utils.hpp index 73996f6..969741f 100644 --- a/src/face_utils.hpp +++ b/src/face_utils.hpp @@ -24,6 +24,28 @@ inline cv::Mat align_face(const cv::Mat& img, return crop; } +// ── enhance_for_retry ──────────────────────────────────────────────────────── +// Used when initial face detection finds nothing. Pads the image by 50% +// (border-replicated, so the detector doesn't see a hard edge) and applies +// CLAHE to boost local contrast, giving the detector a second try. +inline cv::Mat enhance_for_retry(const cv::Mat& img) { + cv::Mat padded; + const int pad_x = img.cols / 4; + const int pad_y = img.rows / 4; + cv::copyMakeBorder(img, padded, pad_y, pad_y, pad_x, pad_x, cv::BORDER_REPLICATE); + + cv::Mat lab; + cv::cvtColor(padded, lab, cv::COLOR_BGR2Lab); + std::vector<cv::Mat> channels; + cv::split(lab, channels); + cv::createCLAHE(2.0, cv::Size(8, 8))->apply(channels[0], channels[0]); + cv::merge(channels, lab); + + cv::Mat out; + cv::cvtColor(lab, out, cv::COLOR_Lab2BGR); + return out; +} + // ── l2_normalise ────────────────────────────────────────────────────────────── inline Embedding l2_normalise(const float* row) { float norm = 0.f;