"""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)