dtourolle/scene-actor-extraction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

faster calibration curve generation

Browse Source 
jellyfin intergration
This commit is contained in:
dtourolle
2026-06-12 17:54:23 +02:00
parent d753062c6c
commit a1d6759abc
17 changed files with 1379 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/make_gallery.py
+27 -71
View File
@@ -1,9 +1,9 @@
#!/usr/bin/env python3
"""make_gallery.py — fetch actor images for a movie and build gallery.json.
Fetches the cast from TMDB, downloads actor profile images, runs the C++
embed_faces binary (SCRFD + ArcFace, same models as scene_analyze) to produce
embeddings, then writes gallery.json.
Fetches the cast from TMDB, downloads actor profile images, embeds them via
the sae_embed module (SCRFD + ArcFace, same models as scene_analyze, loaded
once), then writes gallery.json.
Requirements:
pip install requests Pillow
@@ -22,7 +22,7 @@ Usage:
--output gallery.json
# Additional options:
# --embed-bin build/embed_faces path to embed_faces binary
# --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
@@ -32,19 +32,18 @@ Get a free TMDB API key at: https://www.themoviedb.org/settings/api
"""
import argparse
import io
import json
import os
import subprocess
import sys
import tempfile
import time
from pathlib import Path
import io
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"
@@ -81,7 +80,7 @@ def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]:
# Get IMDB ID for this person
ext = tmdb_get(f"/person/{person_id}/external_ids", key)
imdb_id = ext.get("imdb_id") or f"tmdb_{person_id}"
imdb_id = ext.get("imdb_id") or ""
# Get profile images (sorted by vote_average desc by TMDB)
images_data = tmdb_get(f"/person/{person_id}/images", key)
@@ -96,6 +95,7 @@ def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]:
"id": person_id,
"name": member["name"],
"imdb_id": imdb_id,
"tmdb_id": str(person_id),
"profile_images": image_urls,
})
time.sleep(0.05) # be polite to TMDB
@@ -124,42 +124,9 @@ def download_images(actor: dict, dest_dir: Path, n: int) -> list[Path]:
return paths
# ── Embedding via embed_faces binary ─────────────────────────────────────────
def embed_images(image_paths: list[Path], embed_bin: str,
detector: str, arcface: str) -> list[dict | None]:
"""
Call the C++ embed_faces binary on a list of images.
Returns a list of result dicts (or None if no face / error) per image.
"""
if not image_paths:
return []
cmd = [
embed_bin,
"--detector", detector,
"--arcface", arcface,
] + [str(p) for p in image_paths]
try:
proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
except subprocess.CalledProcessError as e:
print(f"[error] embed_faces failed:\n{e.stderr}", file=sys.stderr)
return [None] * len(image_paths)
try:
results = json.loads(proc.stdout)
except json.JSONDecodeError as e:
print(f"[error] embed_faces output is not valid JSON: {e}", file=sys.stderr)
return [None] * len(image_paths)
return results
# ── Gallery assembly ─────────────────────────────────────────────────────────
def build_gallery(movie_id: int, key: str, embed_bin: str,
detector: str, arcface: str,
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."""
@@ -171,27 +138,27 @@ def build_gallery(movie_id: int, key: str, embed_bin: str,
for actor in actors:
safe_name = actor["name"].replace(" ", "_")
actor_dir = image_root / f"{actor['imdb_id']}_{safe_name}"
dir_id = actor["imdb_id"] or f"tmdb_{actor['tmdb_id']}"
actor_dir = image_root / f"{dir_id}_{safe_name}"
print(f"\n{actor['name']} ({actor['imdb_id']})", file=sys.stderr)
print(f"\n{actor['name']} ({dir_id})", file=sys.stderr)
image_paths = download_images(actor, actor_dir, images_per_actor)
if not image_paths:
print(" no images downloaded, skipping", file=sys.stderr)
continue
print(f" embedding {len(image_paths)} image(s)…", file=sys.stderr)
results = embed_images(image_paths, embed_bin, detector, arcface)
embeddings = []
source_images = []
for path, res in zip(image_paths, results):
if res is None or res.get("embedding") is None:
reason = res.get("error", "unknown") if res else "binary error"
print(f" [skip] {path.name}: {reason}", file=sys.stderr)
for path in image_paths:
res = embedder.embed(str(path))
if not res.ok:
print(f" [skip] {path.name}: {res.error}", file=sys.stderr)
continue
embeddings.append(res["embedding"])
embeddings.append(res.embedding)
source_images.append(path.name)
print(f" [ok] {path.name} conf={res.get('confidence', 0):.2f}",
print(f" [ok] {path.name} conf={res.confidence:.2f}",
file=sys.stderr)
if not embeddings:
@@ -200,6 +167,8 @@ def build_gallery(movie_id: int, key: str, embed_bin: str,
gallery_actors.append({
"imdb_id": actor["imdb_id"],
"tmdb_id": actor["tmdb_id"],
"jellyfin_id": "",
"name": actor["name"],
"source_images": source_images,
"embeddings": embeddings,
@@ -213,7 +182,7 @@ def build_gallery(movie_id: int, key: str, embed_bin: str,
def main():
parser = argparse.ArgumentParser(
description="Fetch TMDB cast images and build gallery.json via embed_faces")
description="Fetch TMDB cast images and build gallery.json via sae_embed")
parser.add_argument("--tmdb-key", required=True,
help="TMDB Bearer token (API Read Access Token from themoviedb.org/settings/api)")
group = parser.add_mutually_exclusive_group(required=True)
@@ -222,8 +191,8 @@ def main():
group.add_argument("--movie-id", type=int,
help="TMDB movie ID (alternative to --imdb-id)")
parser.add_argument("--output", required=True, help="Output gallery.json path")
parser.add_argument("--embed-bin", default="build/embed_faces",
help="Path to embed_faces binary (default: build/embed_faces)")
parser.add_argument("--build-dir", default="build",
help="Build directory containing the sae_embed module (default: build)")
parser.add_argument("--models-dir", default="models",
help="Directory containing ONNX models (default: models/)")
parser.add_argument("--arcface", default=None,
@@ -239,21 +208,10 @@ def main():
args = parser.parse_args()
# Resolve paths
embed_bin = str(Path(args.embed_bin).resolve())
models_dir = Path(args.models_dir)
detector = str(models_dir / "scrfd_500m_bnkps.onnx")
arcface = args.arcface if args.arcface else str(models_dir / "arcface_w600k_r50.onnx")
output = Path(args.output)
image_root = Path(args.image_dir) if args.image_dir else output.parent / "images"
# Validate
if not Path(embed_bin).is_file():
sys.exit(f"embed_faces binary not found: {embed_bin}\n"
f"Build it first: cmake --build build --target embed_faces")
for model, name in [(detector, "SCRFD"), (arcface, "ArcFace")]:
if not Path(model).is_file():
sys.exit(f"{name} model not found: {model}\n"
f"Run: bash scripts/download_models.sh")
embedder = load_embedder(args.build_dir, args.models_dir, args.arcface)
# Resolve movie ID
movie_id = args.movie_id
@@ -266,9 +224,7 @@ def main():
gallery = build_gallery(
movie_id = movie_id,
key = args.tmdb_key,
embed_bin = embed_bin,
detector = detector,
arcface = arcface,
embedder = embedder,
max_actors = args.max_actors,
images_per_actor = args.images_per_actor,
image_root = image_root,