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:
+4
-86
@@ -31,69 +31,15 @@ Get a free TMDB API key at: https://www.themoviedb.org/settings/api
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
from sae_embed_loader import load_embedder
|
||||
|
||||
TMDB_BASE = "https://api.themoviedb.org/3"
|
||||
TMDB_IMG = "https://image.tmdb.org/t/p/original"
|
||||
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)",
|
||||
}
|
||||
|
||||
# ── TMDB helpers ──────────────────────────────────────────────────────────────
|
||||
|
||||
def tmdb_get(path: str, token: str, **params) -> dict:
|
||||
url = TMDB_BASE + path
|
||||
if token.startswith("eyJ"):
|
||||
headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"}
|
||||
r = requests.get(url, params=params, headers=headers, timeout=10)
|
||||
else:
|
||||
params["api_key"] = token
|
||||
r = requests.get(url, params=params, headers={"Accept": "application/json"}, timeout=10)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
# ── Wikidata helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
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 Exception as e:
|
||||
print(f" [warn] wikidata lookup failed for {imdb_person_id}: {e}", file=sys.stderr)
|
||||
return []
|
||||
|
||||
|
||||
def tmdb_id_from_imdb(imdb_id: str, key: str) -> int:
|
||||
data = tmdb_get(f"/find/{imdb_id}", key, external_source="imdb_id")
|
||||
results = data.get("movie_results", [])
|
||||
if not results:
|
||||
raise ValueError(f"No TMDB movie found for IMDB ID {imdb_id}")
|
||||
return results[0]["id"]
|
||||
from sae_gallery import download_images, save_gallery, wikidata_image_urls
|
||||
from sae_tmdb import TMDB_IMG, tmdb_get, tmdb_id_from_imdb
|
||||
|
||||
|
||||
def fetch_cast(movie_id: int, key: str) -> list[dict]:
|
||||
@@ -135,27 +81,6 @@ def fetch_cast(movie_id: int, key: str) -> list[dict]:
|
||||
return actors
|
||||
|
||||
|
||||
# ── Image download ────────────────────────────────────────────────────────────
|
||||
|
||||
def download_images(actor: dict, dest_dir: Path, n: int) -> list[Path]:
|
||||
"""Download up to n profile images for an actor into dest_dir."""
|
||||
dest_dir.mkdir(parents=True, exist_ok=True)
|
||||
paths = []
|
||||
for i, url in enumerate(actor["profile_images"][:n]):
|
||||
out = dest_dir / f"{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] download failed: {url}: {e}", file=sys.stderr)
|
||||
return paths
|
||||
|
||||
|
||||
# ── Gallery assembly ─────────────────────────────────────────────────────────
|
||||
|
||||
def build_gallery(movie_id: int, key: str, embedder,
|
||||
@@ -182,7 +107,7 @@ def build_gallery(movie_id: int, key: str, embedder,
|
||||
"tmdb_id": actor["tmdb_id"], "reason": "no images found"})
|
||||
continue
|
||||
|
||||
image_paths = download_images(actor, actor_dir, images_per_actor)
|
||||
image_paths = download_images(actor["profile_images"], actor_dir, images_per_actor)
|
||||
if not image_paths:
|
||||
print(" no images downloaded, skipping", file=sys.stderr)
|
||||
missing.append({"name": actor["name"], "imdb_id": actor["imdb_id"],
|
||||
@@ -279,14 +204,7 @@ def main():
|
||||
if n_actors == 0:
|
||||
sys.exit("No actors could be processed — check models and images.")
|
||||
|
||||
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