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:
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user