improved jellyfin support

This commit is contained in:
2026-06-12 20:57:33 +02:00
parent a1d6759abc
commit fc16d4a0e1
8 changed files with 335 additions and 65 deletions
+61 -15
View File
@@ -24,7 +24,6 @@ Usage:
# Additional options:
# --build-dir build/ build dir containing sae_embed module
# --models-dir models/ directory with ONNX models
# --max-actors 20 how many cast members to include
# --images-per-actor 3 profile images to download per actor
# --image-dir /tmp/gallery_imgs where to cache downloaded images
@@ -46,6 +45,11 @@ 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 ──────────────────────────────────────────────────────────────
@@ -61,6 +65,29 @@ def tmdb_get(path: str, token: str, **params) -> dict:
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", [])
@@ -69,10 +96,10 @@ def tmdb_id_from_imdb(imdb_id: str, key: str) -> int:
return results[0]["id"]
def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]:
def fetch_cast(movie_id: int, key: str) -> list[dict]:
"""Return list of {id, name, imdb_id, profile_images: [...url...]}."""
credits = tmdb_get(f"/movie/{movie_id}/credits", key)
cast = credits.get("cast", [])[:max_actors]
cast = credits.get("cast", [])
actors = []
for member in cast:
@@ -88,8 +115,13 @@ def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]:
image_urls = [TMDB_IMG + p["file_path"] for p in profiles if p.get("file_path")]
if not image_urls:
print(f" [warn] no images for {member['name']}, skipping", file=sys.stderr)
continue
image_urls = wikidata_image_urls(imdb_id)
if image_urls:
print(f" [info] no TMDB images for {member['name']}, "
f"found {len(image_urls)} via Wikidata", file=sys.stderr)
if not image_urls:
print(f" [warn] no images for {member['name']}", file=sys.stderr)
actors.append({
"id": person_id,
@@ -127,14 +159,15 @@ def download_images(actor: dict, dest_dir: Path, n: int) -> list[Path]:
# ── Gallery assembly ─────────────────────────────────────────────────────────
def build_gallery(movie_id: int, key: str, embedder,
max_actors: int, images_per_actor: int,
image_root: Path) -> dict:
"""Fetch cast, download images, embed, return gallery dict."""
images_per_actor: int,
image_root: Path) -> tuple[dict, list[dict]]:
"""Fetch cast, download images, embed, return (gallery dict, actors needing more images)."""
print(f"Fetching cast for TMDB movie {movie_id}", file=sys.stderr)
actors = fetch_cast(movie_id, key, max_actors)
print(f"Found {len(actors)} actors with images", file=sys.stderr)
actors = fetch_cast(movie_id, key)
print(f"Found {len(actors)} cast member(s)", file=sys.stderr)
gallery_actors = []
missing = []
for actor in actors:
safe_name = actor["name"].replace(" ", "_")
@@ -142,9 +175,18 @@ def build_gallery(movie_id: int, key: str, embedder,
actor_dir = image_root / f"{dir_id}_{safe_name}"
print(f"\n{actor['name']} ({dir_id})", file=sys.stderr)
if not actor["profile_images"]:
print(" no images found, skipping", file=sys.stderr)
missing.append({"name": actor["name"], "imdb_id": actor["imdb_id"],
"tmdb_id": actor["tmdb_id"], "reason": "no images found"})
continue
image_paths = download_images(actor, 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"],
"tmdb_id": actor["tmdb_id"], "reason": "download failed"})
continue
print(f" embedding {len(image_paths)} image(s)…", file=sys.stderr)
@@ -163,6 +205,8 @@ def build_gallery(movie_id: int, key: str, embedder,
if not embeddings:
print(" no valid embeddings, skipping actor", file=sys.stderr)
missing.append({"name": actor["name"], "imdb_id": actor["imdb_id"],
"tmdb_id": actor["tmdb_id"], "reason": "no valid embeddings"})
continue
gallery_actors.append({
@@ -175,7 +219,7 @@ def build_gallery(movie_id: int, key: str, embedder,
})
print(f"{len(embeddings)} embedding(s) stored", file=sys.stderr)
return {"actors": gallery_actors}
return {"actors": gallery_actors}, missing
# ── Entry point ───────────────────────────────────────────────────────────────
@@ -197,8 +241,6 @@ 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("--max-actors", type=int, default=20,
help="Maximum number of cast members to include (default: 20)")
parser.add_argument("--images-per-actor",type=int, default=3,
help="Profile images to download per actor (default: 3)")
parser.add_argument("--image-dir", default=None,
@@ -221,11 +263,10 @@ def main():
print(f"TMDB movie ID: {movie_id}", file=sys.stderr)
# Build gallery
gallery = build_gallery(
gallery, missing = build_gallery(
movie_id = movie_id,
key = args.tmdb_key,
embedder = embedder,
max_actors = args.max_actors,
images_per_actor = args.images_per_actor,
image_root = image_root,
)
@@ -242,6 +283,11 @@ def main():
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)
if __name__ == "__main__":
main()