scene-actor-extraction/scripts/sae_gallery.py

"""Shared image-download and gallery-writing helpers.

Consolidates the three near-identical download loops (make_gallery.download_images,
make_jellyfin_gallery.download_urls + download_person_images) and the duplicated
"write gallery.h5 + .missing_images.json" tail from both builders.

Galleries are written directly as HDF5 — never JSON. Same layout the C++ side
reads/writes (src/gallery/gallery_store.cpp): flat [N,512] embeddings + per-actor
offset/count, parallel imdb_id/tmdb_id/jellyfin_id/name string arrays, a
per-embedding-row source_images array, and an /embedder group carrying the
GR-004 model binding. calibration is left absent (calib_hash=0); the C++
identity_matcher fits and writes it back into the file on first use.

The GR-004 embedder stamp written into that /embedder group lives in sae_stamp
and is re-exported below, so existing callers keep importing it from here.
"""

import io
import json
import sys
from pathlib import Path

import h5py
import numpy as np
import requests
from PIL import Image, UnidentifiedImageError

# Errors that mean "this one image couldn't be fetched/decoded" rather than a bug:
# network/HTTP failures, filesystem errors, and undecodable image bytes. We catch
# exactly these around a single download so a coding error isn't swallowed as a
# spurious "download failed".
DOWNLOAD_ERRORS = (requests.RequestException, OSError, UnidentifiedImageError)

# Wikidata SPARQL: map an IMDB person id (P345) to their image (P18) on
# Wikimedia Commons. Used as a free, CC-licensed headshot fallback when a
# provider (TMDB/Jellyfin) has no usable image. Wikidata requires a
# descriptive User-Agent.
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)",
}


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 (requests.RequestException, ValueError) as e:
        print(f"  [warn] wikidata lookup failed for {imdb_person_id}: {e}", file=sys.stderr)
        return []


def download_image(url: str, out_path: Path, *,
                   headers: dict | None = None,
                   params: dict | None = None) -> Path | None:
    """Download one image to out_path as JPEG, returning the path or None on failure.

    Skips the download if out_path already exists and is larger than 1 KiB
    (a previously cached, non-truncated image).
    """
    if out_path.exists() and out_path.stat().st_size > 1024:
        return out_path
    try:
        r = requests.get(url, headers=headers, params=params, timeout=15)
        r.raise_for_status()
        ctype = r.headers.get("Content-Type", "")
        if ctype and not ctype.startswith("image/"):
            print(f"  [warn] unexpected response for {url}: "
                  f"status={r.status_code} content-type={ctype!r} len={len(r.content)}",
                  file=sys.stderr)
            return None
        out_path.parent.mkdir(parents=True, exist_ok=True)
        Image.open(io.BytesIO(r.content)).convert("RGB").save(out_path, "JPEG")
        return out_path
    except DOWNLOAD_ERRORS as e:
        print(f"  [warn] image download failed: {url}: {e}", file=sys.stderr)
        return None


def download_images(urls: list[str], dest_dir: Path, n: int,
                    start_index: int = 0, *,
                    headers: dict | None = None,
                    params: dict | None = None) -> 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"
        path = download_image(url, out, headers=headers, params=params)
        if path is not None:
            paths.append(path)
    return paths


# ── GR-004: gallery ↔ embedder binding ───────────────────────────────────────
# Implemented in sae_stamp (kept dependency-light so the optimizer's replay
# subprocesses can import it without pulling requests/Pillow); re-exported here
# because the gallery writers and every existing caller reach for it via this
# module. See src/gallery/embedder_stamp.hpp for the C++ twin and the rationale.
from sae_stamp import (  # noqa: F401
    EmbedderMismatch,
    check_embedder_stamp,
    describe_stamp,
    embedder_stamp,
    enforce_embedder_stamp,
    read_gallery_stamp,
    require_gallery_stamp_from_env,
    sha256_file,
    verify_gallery_stamp,
    _as_str,
    _stamp_empty,
)


def save_gallery_hdf5(gallery: dict, output: Path, embedder: dict | None = None) -> None:
    """Write a gallery dict ({"actors": [...]}) directly as HDF5 — same schema
    src/gallery/gallery_store.cpp reads/writes. No calibration group; the
    C++ identity_matcher computes and writes it back into this file on first
    use against an unseen set of embeddings.

    `embedder` is the GR-004 stamp (see embedder_stamp()); it may also be carried
    on the gallery dict under "embedder", which is how a filtered/derived gallery
    keeps its binding without the caller having to re-hash anything."""
    embedder = embedder if embedder is not None else gallery.get("embedder")
    actors = gallery["actors"]
    embs, offsets, counts = [], [], []
    imdb, tmdb, jf, name, src_images = [], [], [], [], []
    row = 0
    for a in actors:
        e = a.get("embeddings", [])
        offsets.append(row)
        counts.append(len(e))
        row += len(e)
        embs.extend(e)
        si = a.get("source_images", [])
        for i in range(len(e)):
            src_images.append(si[i] if i < len(si) else "")
        imdb.append(a.get("imdb_id", "") or "")
        tmdb.append(str(a.get("tmdb_id", "") or ""))
        jf.append(a.get("jellyfin_id", a.get("jellyfin_person_id", "")) or "")
        name.append(a.get("name", "") or "")

    emb_arr = np.asarray(embs, dtype=np.float32) if embs else np.zeros((0, 512), np.float32)
    if emb_arr.ndim == 1:
        emb_arr = emb_arr.reshape(0, 512)

    output.parent.mkdir(parents=True, exist_ok=True)
    str_t = h5py.string_dtype("utf-8")
    with h5py.File(output, "w") as f:
        f.create_dataset("embeddings", data=emb_arr)
        f.create_dataset("offset", data=np.asarray(offsets, np.int64))
        f.create_dataset("count", data=np.asarray(counts, np.int32))
        f.create_dataset("imdb_id", data=np.asarray(imdb, dtype=object), dtype=str_t)
        f.create_dataset("tmdb_id", data=np.asarray(tmdb, dtype=object), dtype=str_t)
        f.create_dataset("jellyfin_id", data=np.asarray(jf, dtype=object), dtype=str_t)
        f.create_dataset("name", data=np.asarray(name, dtype=object), dtype=str_t)
        f.create_dataset("source_images", data=np.asarray(src_images, dtype=object), dtype=str_t)
        # TRACES: GR-004 | SR-001 — omitted entirely when unknown, so "unstamped"
        # round-trips as unstamped rather than as a stamp naming no model.
        if not _stamp_empty(embedder):
            g = f.create_group("embedder")
            g.attrs["model_name"] = embedder.get("model_name", "")
            g.attrs["model_sha256"] = embedder.get("model_sha256", "")
            g.attrs["embed_dim"] = np.int32(embedder.get("embed_dim", 512))
    stamp_note = (f", embedder {embedder['model_name']}" if not _stamp_empty(embedder)
                  else ", NO EMBEDDER STAMP (GR-004)")
    print(f"Saved: {output} ({len(actors)} actors, {emb_arr.shape[0]} embeddings"
          f"{stamp_note})", file=sys.stderr)


def load_gallery_hdf5(path: Path) -> dict:
    """Read an HDF5 gallery back into the same {"actors": [...]} dict shape the
    builders work with in memory (for --merge). Mirrors save_gallery_hdf5."""
    with h5py.File(path, "r") as f:
        emb = f["embeddings"][:]
        offset = f["offset"][:]
        count = f["count"][:]
        imdb = [s.decode() if isinstance(s, bytes) else s for s in f["imdb_id"][:]]
        tmdb = [s.decode() if isinstance(s, bytes) else s for s in f["tmdb_id"][:]]
        jf = [s.decode() if isinstance(s, bytes) else s for s in f["jellyfin_id"][:]]
        name = [s.decode() if isinstance(s, bytes) else s for s in f["name"][:]]
        src_images = None
        if "source_images" in f:
            src_images = [s.decode() if isinstance(s, bytes) else s
                         for s in f["source_images"][:]]
        # TRACES: GR-004 | SR-001 — carried through so a derived gallery (filter,
        # merge, cast-restrict) keeps the binding of the gallery it came from.
        stamp = None
        if "embedder" in f:
            a = f["embedder"].attrs
            stamp = {"model_name": _as_str(a.get("model_name", "")),
                     "model_sha256": _as_str(a.get("model_sha256", "")),
                     "embed_dim": int(a.get("embed_dim", 512))}

    actors = []
    for a in range(len(offset)):
        s, n = int(offset[a]), int(count[a])
        actor = {"imdb_id": imdb[a], "tmdb_id": tmdb[a], "jellyfin_id": jf[a],
                "name": name[a], "embeddings": [emb[s + i].tolist() for i in range(n)]}
        if src_images is not None:
            actor["source_images"] = [src_images[s + i] for i in range(n)]
        actors.append(actor)
    out = {"actors": actors}
    if stamp is not None:
        out["embedder"] = stamp
    return out


def save_gallery(gallery: dict, missing: list[dict], output: Path,
                 embedder: dict | None = None) -> None:
    """Write the gallery as HDF5 (forcing a .h5 extension) and, if any actors
    lack images, a .missing_images.json sidecar."""
    if output.suffix not in (".h5", ".hdf5"):
        output = output.with_suffix(".h5")
    save_gallery_hdf5(gallery, output, embedder)

    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)