scene-actor-extraction/scripts/optimizer/reembed_gallery.py

#!/usr/bin/env python3
"""
reembed_gallery.py — re-embed an existing gallery's actors with a different model.

For the embedding-model bake-off: take a reference gallery (with all actor ids +
source_images) and produce a new gallery where every actor's embeddings are computed
by a DIFFERENT ArcFace/LVFace model from the SAME cached source images. All identity
keys (imdb/tmdb/jellyfin/name) are preserved, so membership/matching is unchanged —
only the embedding vectors (and hence the model's similarity space) differ.

Source images live in `--images <root>/<jellyfin_id>_<Name>/NN.jpg` (the gallery build
cache). Actors are matched to their image dir by jellyfin_id first, then name.

Usage:
  python scripts/optimizer/reembed_gallery.py \
      --ref gallery_arcface_w600k_r50.h5 \
      --images images \
      --arcface models/arcface_r18.onnx \
      --out experiments/galleries/gallery_arcface_r18.h5 \
      [--build-dir build]
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts"))
from sae_embed_loader import load_embedder, resolve_arcface  # noqa: E402
from sae_gallery import (embedder_stamp, load_gallery_hdf5,  # noqa: E402
                         save_gallery_hdf5)


def find_dir(images_root: Path, jellyfin_id: str, name: str) -> Path | None:
    if jellyfin_id:
        d = images_root / f"{jellyfin_id}_{name.replace(' ', '_')}"
        if d.is_dir():
            return d
        # jellyfin_id prefix match (name spelling may differ)
        hits = list(images_root.glob(f"{jellyfin_id}_*"))
        if hits:
            return hits[0]
    hits = list(images_root.glob(f"*_{name.replace(' ', '_')}"))
    return hits[0] if hits else None


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--ref", required=True, help="reference gallery.h5 (ids + source imgs)")
    p.add_argument("--images", required=True, help="image cache root")
    p.add_argument("--arcface", required=True, help="model ONNX to re-embed with")
    p.add_argument("--out", required=True)
    p.add_argument("--build-dir", default=str(REPO / "build"))
    p.add_argument("--models-dir", default=str(REPO / "models"))
    args = p.parse_args()

    ref = load_gallery_hdf5(Path(args.ref))
    images_root = Path(args.images)
    embedder = load_embedder(args.build_dir, args.models_dir, args.arcface)
    # TRACES: GR-004 | SR-001 — this script exists to produce a gallery in a
    # DIFFERENT model's space from the reference. The output must therefore never
    # inherit the reference's stamp; it carries the stamp of --arcface, which is
    # the whole point of the bake-off being safe to run.
    stamp = embedder_stamp(resolve_arcface(args.models_dir, args.arcface))

    out_actors = []
    n_ok = n_nodir = n_noemb = 0
    total = len(ref["actors"])
    for i, a in enumerate(ref["actors"], 1):
        d = find_dir(images_root, a.get("jellyfin_id", ""), a["name"])
        if d is None:
            n_nodir += 1
            continue
        embeddings = []
        for img in sorted(d.glob("*.jpg")):
            res = embedder.embed(str(img))
            if res.ok:
                embeddings.append(list(res.embedding))
        if not embeddings:
            n_noemb += 1
            continue
        out_actors.append({"imdb_id": a.get("imdb_id", ""), "tmdb_id": a.get("tmdb_id", ""),
                           "jellyfin_id": a.get("jellyfin_id", ""), "name": a["name"],
                           "embeddings": embeddings,
                           "source_images": [p.name for p in sorted(d.glob("*.jpg"))]})
        n_ok += 1
        if i % 200 == 0 or i == total:
            print(f"  [{i}/{total}] ok={n_ok} no_dir={n_nodir} no_emb={n_noemb}",
                  file=sys.stderr)

    save_gallery_hdf5({"actors": out_actors}, Path(args.out), stamp)
    n_emb = sum(len(a["embeddings"]) for a in out_actors)
    print(f"[reembed] {Path(args.arcface).stem}: {n_ok}/{total} actors, {n_emb} embeddings "
          f"→ {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()