improved jellyfin support
This commit is contained in:
+61
-15
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user