From 458116f118e8e27798319e80a6fc6b06550872ee Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Thu, 30 Jul 2026 17:32:08 +0200 Subject: [PATCH] fix(trt): drop explicit shapes for static TransNetV2; gallery over-fetch + dedup trtexec rejects --minShapes/--optShapes/--maxShapes for a fully static model ("Static model does not take explicit shapes"). TransNetV2's input is fixed at 1x100x27x48x3, so the shape comes from the model itself. Gallery build now over-fetches TMDB/Wikidata candidates by a configurable factor: near-duplicate stills (the same photo at different crops or resolutions) are discarded after embedding, so downloading exactly images_per_actor left actors short of that many *distinct* embeddings. Co-Authored-By: Claude Opus 5 --- scripts/build_trt_engines.sh | 10 ++-- scripts/make_jellyfin_gallery.py | 98 +++++++++++++++++++++++++++----- scripts/sae_embed_loader.py | 12 +++- src/face_embedder_engine.hpp | 17 ++++-- src/python_bindings.cpp | 6 +- 5 files changed, 117 insertions(+), 26 deletions(-) diff --git a/scripts/build_trt_engines.sh b/scripts/build_trt_engines.sh index 0c7dae1..4a8d0b6 100755 --- a/scripts/build_trt_engines.sh +++ b/scripts/build_trt_engines.sh @@ -76,14 +76,14 @@ if [[ -f "$SCENE_MODEL" ]]; then echo "== TransNetV2 (scene detector) ==" # Fixed 1x100x27x48x3 window. The raw-TRT scene detector backend loads this # engine directly via --scene-detector-engine; the ORT-TRT EP builds its own. - SCENE_IN="$(input_name "$SCENE_MODEL")" - echo " (input tensor: $SCENE_IN)" + # No --*Shapes here: TransNetV2's input is fully static (1x100x27x48x3 + # with no dynamic dimensions), and TensorRT rejects explicit shape + # profiles for such a model — "Static model does not take explicit shapes + # since the shape of inference tensors will be determined by the model + # itself". The shape comes from the model. run trtexec \ --onnx="$SCENE_MODEL" \ --fp16 \ - --minShapes="$SCENE_IN":1x100x27x48x3 \ - --optShapes="$SCENE_IN":1x100x27x48x3 \ - --maxShapes="$SCENE_IN":1x100x27x48x3 \ --saveEngine="$OUT/transnetv2.100x27x48.fp16.engine" \ --useCudaGraph else diff --git a/scripts/make_jellyfin_gallery.py b/scripts/make_jellyfin_gallery.py index 0ff4153..45379cd 100644 --- a/scripts/make_jellyfin_gallery.py +++ b/scripts/make_jellyfin_gallery.py @@ -147,11 +147,14 @@ def download_person_images(base_url: str, api_key: str, person_id: str, 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 + fetch_imdb_ids: bool, tmdb_key: str | None, + fetch_overfetch: float = 1.0 ) -> 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) + # Over-fetch target: see the TMDB block below. + tmdb_budget = int(images_per_actor * fetch_overfetch) image_paths = download_person_images(base_url, api_key, pid, actor_dir, images_per_actor) # Need the IMDB id for the TMDB /find lookup, to persist it (--fetch-imdb-ids), @@ -177,8 +180,13 @@ def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict, 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) + # Over-fetch from TMDB: near-duplicate stills (the same photo at different + # crops/resolutions) are dropped after embedding, so downloading exactly + # images_per_actor would leave the actor short of that many *distinct* + # embeddings. Pulling extra candidates lets the dedup filter discard + # duplicates while still reaching the target. + if len(image_paths) < tmdb_budget and tmdb_urls: + needed = tmdb_budget - 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_images(tmdb_urls, actor_dir, needed, start_index=len(image_paths)) @@ -186,10 +194,10 @@ def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict, # Last resort: a CC-licensed Commons headshot via Wikidata, keyed by the # actor's IMDB id. Catches actors TMDB has no usable image for (or that the # name search missed entirely). - if len(image_paths) < images_per_actor and imdb_id: + if len(image_paths) < tmdb_budget and imdb_id: wiki_urls = wikidata_image_urls(imdb_id) if wiki_urls: - needed = images_per_actor - len(image_paths) + needed = tmdb_budget - len(image_paths) print(f" {name}: still short, falling back to Wikidata/Commons " f"({len(wiki_urls)} image(s) available)…", file=sys.stderr) image_paths += download_images(wiki_urls, actor_dir, needed, @@ -198,9 +206,39 @@ def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict, return image_paths, imdb_id, tmdb_id +# Default cosine-distance tolerance below which two embeddings of the same +# actor are treated as the same image. Embeddings are L2-normalised by the +# backend, so cosine similarity is a plain dot product and the distance is +# 1 - dot. Expanding an actor's photo set via TMDB frequently returns the same +# still at different crops/resolutions; those embed to nearly identical vectors +# and add gallery size and match cost without adding information. +DEDUP_TOL = 1e-3 + + +def _cosine(a, b) -> float: + """Cosine similarity of two L2-normalised embeddings.""" + return float(sum(x * y for x, y in zip(a, b))) + + +def _near_duplicate(emb, existing, tol: float) -> int | None: + """Index of the first embedding within `tol` cosine distance of `emb`. + + Returns None when `emb` is sufficiently distinct from everything in + `existing`. tol <= 0 disables the check. + """ + if tol <= 0: + return None + for i, prev in enumerate(existing): + if 1.0 - _cosine(emb, prev) < tol: + return i + return None + + 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) -> tuple[dict | None, str | None]: + embed_executor: concurrent.futures.ThreadPoolExecutor, + dedup_tol: float = DEDUP_TOL, + max_embeddings: int = 0) -> 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 @@ -217,20 +255,35 @@ def embed_actor(pid: str, info: dict, image_paths: list[Path], embeddings = [] source_images = [] + n_dup = 0 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) + emb = res.embedding + dup = _near_duplicate(emb, embeddings, dedup_tol) + if dup is not None: + n_dup += 1 + print(f" [dup] {name}/{path.name}: matches {source_images[dup]} " + f"(cos={_cosine(emb, embeddings[dup]):.6f}), not stored", + file=sys.stderr) + continue + embeddings.append(emb) source_images.append(path.name) + # Stop once we have the requested number of *distinct* embeddings; the + # extra candidates were only fetched to absorb duplicates. + if max_embeddings and len(embeddings) >= max_embeddings: + break if not embeddings: print(f" {name}: no valid embeddings, skipping actor", file=sys.stderr) return None, "no valid embeddings" - print(f" {name}: → {len(embeddings)} embedding(s) stored", file=sys.stderr) + dup_note = f" ({n_dup} near-duplicate(s) dropped)" if n_dup else "" + print(f" {name}: → {len(embeddings)} embedding(s) stored{dup_note}", + file=sys.stderr) return { "imdb_id": imdb_id if (fetch_imdb_ids and imdb_id) else "", "tmdb_id": tmdb_id or "", @@ -245,12 +298,16 @@ def embed_actor(pid: str, info: dict, image_paths: list[Path], 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) -> tuple[dict | None, str | None]: + embed_executor: concurrent.futures.ThreadPoolExecutor, + dedup_tol: float = DEDUP_TOL, + fetch_overfetch: float = 1.0) -> 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( - 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) + base_url, api_key, pid, info, images_per_actor, actor_dir, fetch_imdb_ids, + tmdb_key, fetch_overfetch) + return embed_actor(pid, info, image_paths, imdb_id, tmdb_id, embedder, fetch_imdb_ids, + embed_executor, dedup_tol, images_per_actor) # ── Gallery assembly ───────────────────────────────────────────────────────── @@ -258,7 +315,9 @@ 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) -> tuple[dict, list[dict]]: + tmdb_key: str | None = None, workers: int = 8, + dedup_tol: float = DEDUP_TOL, + fetch_overfetch: float = 1.0) -> tuple[dict, list[dict]]: actors = collect_actors(base_url, api_key, item_types) gallery_actors = [] @@ -298,7 +357,8 @@ def build_gallery(base_url: str, api_key: str, embedder, item_types: list[str], futures = { executor.submit(process_actor, pid, info, base_url, api_key, embedder, images_per_actor, image_root, - fetch_imdb_ids, tmdb_key, embed_executor): (pid, info["name"]) + fetch_imdb_ids, tmdb_key, embed_executor, + dedup_tol, fetch_overfetch): (pid, info["name"]) for pid, info in todo } for future in concurrent.futures.as_completed(futures): @@ -343,6 +403,16 @@ 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("--dedup-tol", type=float, default=DEDUP_TOL, + help=f"Cosine-distance threshold below which a new embedding is treated " + f"as a duplicate of one already stored for that actor and dropped " + f"(default: {DEDUP_TOL}). TMDB often returns the same still at " + f"different crops. Set 0 to keep every embedding.") + parser.add_argument("--overfetch", type=float, default=2.0, + help="Download this multiple of --images-per-actor as candidates, then " + "keep the first N that survive dedup (default: 2.0). Raise it for " + "actors whose TMDB galleries are mostly duplicates; 1.0 disables " + "over-fetching.") parser.add_argument("--images-per-actor", type=int, default=10, help="Images to download per actor (default: 10). Jellyfin usually has " "only 1, so the rest come from the TMDB/Wikidata fallbacks; more " @@ -392,6 +462,8 @@ def main(): image_root=image_root, fetch_imdb_ids=args.fetch_imdb_ids, existing_actors=existing_actors, + dedup_tol=args.dedup_tol, + fetch_overfetch=args.overfetch, tmdb_key=args.tmdb_key, workers=args.workers, ) diff --git a/scripts/sae_embed_loader.py b/scripts/sae_embed_loader.py index 364be83..eb9c001 100644 --- a/scripts/sae_embed_loader.py +++ b/scripts/sae_embed_loader.py @@ -33,4 +33,14 @@ def load_embedder(build_dir: str, models_dir: str, arcface: str | None = None, 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) + # A TRT-backend build cannot load .onnx; it needs pre-built engines from + # scripts/build_trt_engines.sh. Pass them when present (ignored by ORT). + trt = Path(models_path).parent / "trt_cache" + det_engine = trt / "scrfd.scrfd_500m_bnkps.640.fp16.engine" + arc_engine = trt / f"arcface.{Path(arcface_path).stem}.b4.fp16.engine" + + return sae_embed.FaceEmbedder( + detector_path, arcface_path, conf, nms, max_side, + str(det_engine) if det_engine.is_file() else "", + str(arc_engine) if arc_engine.is_file() else "", + ) diff --git a/src/face_embedder_engine.hpp b/src/face_embedder_engine.hpp index a583657..842f4d8 100644 --- a/src/face_embedder_engine.hpp +++ b/src/face_embedder_engine.hpp @@ -32,16 +32,23 @@ struct FaceEmbedResult { class FaceEmbedderEngine { public: + // detector_engine/arcface_engine are optional paths to pre-built TensorRT + // engines. They are required when built with SAE_INFERENCE_BACKEND=TRT + // (which cannot load .onnx directly) and ignored by the ORT backend. FaceEmbedderEngine(const std::string& detector_model, const std::string& arcface_model, - float conf = 0.5f, float nms = 0.4f, int max_side = 500) + float conf = 0.5f, float nms = 0.4f, int max_side = 500, + const std::string& detector_engine = "", + const std::string& arcface_engine = "") : max_side_(max_side) { Config cfg; - cfg.detector_model = detector_model; - cfg.arcface_model = arcface_model; - cfg.detector_conf = conf; - cfg.detector_nms = nms; + cfg.detector_model = detector_model; + cfg.arcface_model = arcface_model; + cfg.detector_engine = detector_engine; + cfg.arcface_engine = arcface_engine; + cfg.detector_conf = conf; + cfg.detector_nms = nms; detector_ = make_face_detector(cfg); embedder_ = make_face_embedder(cfg); } diff --git a/src/python_bindings.cpp b/src/python_bindings.cpp index a9623e5..3356601 100644 --- a/src/python_bindings.cpp +++ b/src/python_bindings.cpp @@ -30,9 +30,11 @@ NB_MODULE(sae_embed, m) { }); nb::class_(m, "FaceEmbedder") - .def(nb::init(), + .def(nb::init(), "detector_model"_a, "arcface_model"_a, - "conf"_a = 0.5f, "nms"_a = 0.4f, "max_side"_a = 500) + "conf"_a = 0.5f, "nms"_a = 0.4f, "max_side"_a = 500, + "detector_engine"_a = "", "arcface_engine"_a = "") .def("embed", &FaceEmbedderEngine::embed_path, "path"_a, nb::call_guard(), "Detect the highest-confidence face in the image, align it, and "