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, and a
per-embedding-row source_images array. calibration is left absent (calib_hash=0);
the C++ identity_matcher fits and writes it back into the file on first use.
"""

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


def save_gallery_hdf5(gallery: dict, output: Path) -> 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."""
    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)
    print(f"Saved: {output} ({len(actors)} actors, {emb_arr.shape[0]} embeddings)",
          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"][:]]

    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)
    return {"actors": actors}


def save_gallery(gallery: dict, missing: list[dict], output: Path) -> 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)

    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)