refactor(scripts): extract shared sae_* helpers and dedupe gallery builders

Consolidate copy-pasted logic across the gallery/run scripts into shared
modules:
  - sae_env.py     — zero-dependency .env loader (populates os.environ)
  - sae_tmdb.py    — TMDB API helpers (tmdb_get, person images, id lookups)
  - sae_jellyfin.py— Jellyfin API helpers (jf_get, id/URL normalisation)
  - sae_gallery.py — image download + gallery.json writing

make_gallery, make_jellyfin_gallery and filter_gallery now import these
instead of carrying their own near-identical copies.
This commit is contained in:
2026-07-04 18:55:37 +02:00
parent aacaefb3dc
commit 152c34b1f4
7 changed files with 431 additions and 258 deletions
+56 -118
View File
@@ -43,31 +43,25 @@ Get a free TMDB API key at: https://www.themoviedb.org/settings/api
import argparse
import concurrent.futures
import os
import io
import json
import sys
from pathlib import Path
import requests
from PIL import Image
sys.path.insert(0, str(Path(__file__).resolve().parent))
import sae_env # noqa: F401 — loads .env into os.environ on import
from sae_embed_loader import load_embedder
TMDB_BASE = "https://api.themoviedb.org/3"
TMDB_IMG = "https://image.tmdb.org/t/p/original"
from sae_gallery import download_image, download_images, save_gallery, wikidata_image_urls
from sae_jellyfin import actor_jellyfin_id, jf_get, normalize_jellyfin_url
from sae_tmdb import (
tmdb_person_by_name,
tmdb_person_for_imdb,
)
# ── Jellyfin API helpers ────────────────────────────────────────────────────
def jf_get(base_url: str, api_key: str, path: str, **params) -> dict:
url = base_url.rstrip("/") + path
headers = {"X-Emby-Token": api_key, "Accept": "application/json"}
r = requests.get(url, params=params, headers=headers, timeout=30)
r.raise_for_status()
return r.json()
def fetch_library_items(base_url: str, api_key: str, item_types: list[str]):
"""Yield every Movie/Series item dict (with its People field)."""
start = 0
@@ -125,99 +119,26 @@ def fetch_imdb_id(base_url: str, api_key: str, person_id: str) -> str | None:
return data.get("ProviderIds", {}).get("Imdb")
# ── TMDB fallback (for actors with no usable Jellyfin image) ────────────────
def tmdb_get(path: str, key: str, **params) -> dict:
url = TMDB_BASE + path
if key.startswith("eyJ"):
headers = {"Authorization": f"Bearer {key}", "Accept": "application/json"}
r = requests.get(url, params=params, headers=headers, timeout=10)
else:
params["api_key"] = key
r = requests.get(url, params=params, headers={"Accept": "application/json"}, timeout=10)
r.raise_for_status()
return r.json()
def tmdb_person_images(tmdb_person_id: str, tmdb_key: str) -> list[str]:
images_data = tmdb_get(f"/person/{tmdb_person_id}/images", tmdb_key)
return [TMDB_IMG + p["file_path"] for p in images_data.get("profiles", []) if p.get("file_path")]
def tmdb_person_for_imdb(imdb_id: str, tmdb_key: str) -> tuple[str | None, list[str]]:
"""Return (tmdb_person_id, profile_image_urls) for the TMDB person matching this IMDB person id."""
data = tmdb_get(f"/find/{imdb_id}", tmdb_key, external_source="imdb_id")
people = data.get("person_results", [])
if not people:
return None, []
tmdb_person_id = str(people[0]["id"])
return tmdb_person_id, tmdb_person_images(tmdb_person_id, tmdb_key)
def tmdb_person_by_name(name: str, tmdb_key: str) -> tuple[str | None, list[str]]:
"""Return (tmdb_person_id, profile_image_urls) for the best name match on TMDB.
Used when Jellyfin has no IMDB ProviderId for this person (the common case —
Jellyfin rarely populates ProviderIds on Person items), so /find/{imdb_id}
isn't an option. /search/person is sorted by popularity; take the top hit.
"""
data = tmdb_get("/search/person", tmdb_key, query=name)
people = data.get("results", [])
if not people:
return None, []
tmdb_person_id = str(people[0]["id"])
return tmdb_person_id, tmdb_person_images(tmdb_person_id, tmdb_key)
def download_urls(urls: list[str], dest_dir: Path, n: int, start_index: int = 0) -> 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"
if out.exists() and out.stat().st_size > 1024:
paths.append(out)
continue
try:
r = requests.get(url, timeout=15)
r.raise_for_status()
Image.open(io.BytesIO(r.content)).convert("RGB").save(out, "JPEG")
paths.append(out)
except Exception as e:
print(f" [warn] TMDB image download failed: {url}: {e}", file=sys.stderr)
return paths
# ── Image download ──────────────────────────────────────────────────────────
def download_person_images(base_url: str, api_key: str, person_id: str,
dest_dir: Path, n: int) -> list[Path]:
"""Download up to n images for a Jellyfin person item into dest_dir."""
"""Download up to n images for a Jellyfin person item into dest_dir.
Jellyfin's per-index Primary image endpoint 404s once you run past the
number of images it has, so stop at the first index that fails rather than
probing all n.
"""
dest_dir.mkdir(parents=True, exist_ok=True)
headers = {"X-Emby-Token": api_key}
paths = []
for i in range(n):
out = dest_dir / f"{i:02d}.jpg"
if out.exists() and out.stat().st_size > 1024:
paths.append(out)
continue
url = base_url.rstrip("/") + f"/Items/{person_id}/Images/Primary/{i}"
try:
r = requests.get(url, headers=headers, params={"api_key": api_key}, timeout=15)
if r.status_code == 404:
break
r.raise_for_status()
ctype = r.headers.get("Content-Type", "")
if not ctype.startswith("image/"):
print(f" [warn] unexpected response for {person_id} index {i}: "
f"status={r.status_code} content-type={ctype!r} "
f"len={len(r.content)} body={r.content[:200]!r}", file=sys.stderr)
break
Image.open(io.BytesIO(r.content)).convert("RGB").save(out, "JPEG")
paths.append(out)
except Exception as e:
print(f" [warn] image download failed for {person_id} index {i}: {e}", file=sys.stderr)
path = download_image(url, out, headers=headers, params={"api_key": api_key})
if path is None:
break
paths.append(path)
return paths
@@ -232,8 +153,12 @@ def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict,
print(f"{name} ({pid}) — in {len(info['appearances'])} title(s)", file=sys.stderr)
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),
# and for the Wikidata fallback below. Fetch it once if Jellyfin's own image
# came up short (so Wikidata can fire) or whenever it's otherwise wanted.
short = len(image_paths) < images_per_actor
imdb_id = None
if fetch_imdb_ids or tmdb_key:
if fetch_imdb_ids or tmdb_key or short:
imdb_id = fetch_imdb_id(base_url, api_key, pid)
# Resolve the TMDB person id whenever possible — independent of whether
@@ -255,7 +180,19 @@ def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict,
needed = images_per_actor - 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_urls(tmdb_urls, actor_dir, needed, start_index=len(image_paths))
image_paths += download_images(tmdb_urls, actor_dir, needed, start_index=len(image_paths))
# 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:
wiki_urls = wikidata_image_urls(imdb_id)
if wiki_urls:
needed = images_per_actor - 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,
start_index=len(image_paths))
return image_paths, imdb_id, tmdb_id
@@ -326,17 +263,29 @@ def build_gallery(base_url: str, api_key: str, embedder, item_types: list[str],
gallery_actors = []
todo = []
n_retry = 0
n_topup = 0
for pid, info in actors.items():
existing = existing_actors.get(pid)
if existing is not None and existing.get("embeddings"):
n_existing = len(existing.get("embeddings", [])) if existing is not None else 0
# Keep an existing actor as-is only once they're at the image target.
# Fewer than that (including zero) → re-process to top up / retry, so a
# merge run can upgrade 1-image actors to images_per_actor without
# re-embedding actors already at the target.
if n_existing >= images_per_actor:
gallery_actors.append(existing)
else:
if existing is not None:
if n_existing > 0:
n_topup += 1
elif existing is not None:
n_retry += 1
todo.append((pid, info))
if len(gallery_actors):
print(f"Skipping {len(gallery_actors)} actor(s) already present in existing gallery", file=sys.stderr)
print(f"Skipping {len(gallery_actors)} actor(s) already at {images_per_actor} "
f"image(s) in existing gallery", file=sys.stderr)
if n_topup:
print(f"Topping up {n_topup} actor(s) with fewer than {images_per_actor} "
f"embedding(s)", file=sys.stderr)
if n_retry:
print(f"Retrying {n_retry} actor(s) with no embeddings in existing gallery", file=sys.stderr)
print(f"Processing {len(todo)} new actor(s) with {workers} worker(s)…", file=sys.stderr)
@@ -393,8 +342,10 @@ 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("--images-per-actor", type=int, default=1,
help="Images to download per actor (default: 1 Jellyfin usually caches one)")
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 "
"images per actor means more embeddings and more robust matching.")
parser.add_argument("--image-dir", default=None,
help="Where to store downloaded images (default: <output_dir>/images)")
parser.add_argument("--fetch-imdb-ids", action="store_true",
@@ -414,12 +365,7 @@ def main():
"on a single dedicated thread regardless of this value.")
args = parser.parse_args()
# Defend against a base URL that accidentally includes an API path,
# e.g. "https://host/Items?" — strip any query string and trailing
# /Items so we don't build double-nested, query-mangled URLs.
jellyfin_url = args.jellyfin_url.split("?", 1)[0].rstrip("/")
if jellyfin_url.endswith("/Items"):
jellyfin_url = jellyfin_url[: -len("/Items")]
jellyfin_url = normalize_jellyfin_url(args.jellyfin_url)
output = Path(args.output)
image_root = Path(args.image_dir) if args.image_dir else output.parent / "images"
@@ -431,8 +377,7 @@ def main():
if args.merge and output.is_file():
existing = json.loads(output.read_text())
for actor in existing.get("actors", []):
# older galleries used "jellyfin_person_id"
pid = actor.get("jellyfin_id") or actor.get("jellyfin_person_id")
pid = actor_jellyfin_id(actor)
if pid:
existing_actors[pid] = actor
print(f"Loaded {len(existing_actors)} actor(s) from existing gallery for merge", file=sys.stderr)
@@ -457,14 +402,7 @@ def main():
if n_actors == 0:
sys.exit("No actors could be processed — check Jellyfin URL/API key and models.")
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(gallery, indent=2) + "\n")
print(f"Saved: {output}", file=sys.stderr)
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)
save_gallery(gallery, missing, output)
if __name__ == "__main__":