Initial commit: scene-actor-extraction pipeline
Source (KPN++ pipeline nodes, ArcFace embedders, SCRFD/YuNet detectors, gallery builder), build scripts, and eval artifacts. - external/KPN as a git submodule (gitea.tourolle.paris/dtourolle/KPN) - ONNX models tracked via Git LFS (models/*.onnx) - generated outputs, TensorRT engines, reference repos, and media ignored
This commit is contained in:
Executable
+291
@@ -0,0 +1,291 @@
|
||||
#!/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.
|
||||
|
||||
Requirements:
|
||||
pip install requests Pillow
|
||||
|
||||
Usage:
|
||||
# By IMDB movie ID (most natural — resolves to TMDB automatically):
|
||||
python scripts/make_gallery.py \\
|
||||
--tmdb-key YOUR_KEY \\
|
||||
--imdb-id tt0137523 \\
|
||||
--output gallery.json
|
||||
|
||||
# Or directly with a TMDB movie ID:
|
||||
python scripts/make_gallery.py \\
|
||||
--tmdb-key YOUR_KEY \\
|
||||
--movie-id 550 \\
|
||||
--output gallery.json
|
||||
|
||||
# Additional options:
|
||||
# --embed-bin build/embed_faces path to embed_faces binary
|
||||
# --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
|
||||
|
||||
Get a free TMDB API key at: https://www.themoviedb.org/settings/api
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import io
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
TMDB_BASE = "https://api.themoviedb.org/3"
|
||||
TMDB_IMG = "https://image.tmdb.org/t/p/original"
|
||||
|
||||
# ── 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()
|
||||
|
||||
|
||||
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"]
|
||||
|
||||
|
||||
def fetch_cast(movie_id: int, key: str, max_actors: int) -> 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]
|
||||
|
||||
actors = []
|
||||
for member in cast:
|
||||
person_id = member["id"]
|
||||
|
||||
# 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}"
|
||||
|
||||
# Get profile images (sorted by vote_average desc by TMDB)
|
||||
images_data = tmdb_get(f"/person/{person_id}/images", key)
|
||||
profiles = images_data.get("profiles", [])
|
||||
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
|
||||
|
||||
actors.append({
|
||||
"id": person_id,
|
||||
"name": member["name"],
|
||||
"imdb_id": imdb_id,
|
||||
"profile_images": image_urls,
|
||||
})
|
||||
time.sleep(0.05) # be polite to TMDB
|
||||
|
||||
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
|
||||
|
||||
|
||||
# ── 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,
|
||||
max_actors: int, images_per_actor: int,
|
||||
image_root: Path) -> dict:
|
||||
"""Fetch cast, download images, embed, return gallery dict."""
|
||||
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)
|
||||
|
||||
gallery_actors = []
|
||||
|
||||
for actor in actors:
|
||||
safe_name = actor["name"].replace(" ", "_")
|
||||
actor_dir = image_root / f"{actor['imdb_id']}_{safe_name}"
|
||||
|
||||
print(f"\n{actor['name']} ({actor['imdb_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)
|
||||
continue
|
||||
embeddings.append(res["embedding"])
|
||||
source_images.append(path.name)
|
||||
print(f" [ok] {path.name} conf={res.get('confidence', 0):.2f}",
|
||||
file=sys.stderr)
|
||||
|
||||
if not embeddings:
|
||||
print(" no valid embeddings, skipping actor", file=sys.stderr)
|
||||
continue
|
||||
|
||||
gallery_actors.append({
|
||||
"imdb_id": actor["imdb_id"],
|
||||
"name": actor["name"],
|
||||
"source_images": source_images,
|
||||
"embeddings": embeddings,
|
||||
})
|
||||
print(f" → {len(embeddings)} embedding(s) stored", file=sys.stderr)
|
||||
|
||||
return {"actors": gallery_actors}
|
||||
|
||||
|
||||
# ── Entry point ───────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Fetch TMDB cast images and build gallery.json via embed_faces")
|
||||
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)
|
||||
group.add_argument("--imdb-id",
|
||||
help="IMDB movie ID, e.g. tt0137523 — looked up via TMDB automatically")
|
||||
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("--models-dir", default="models",
|
||||
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,
|
||||
help="Where to store downloaded images (default: <output_dir>/images)")
|
||||
parser.add_argument("--keep-images", action="store_true",
|
||||
help="Do not delete downloaded images after embedding")
|
||||
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")
|
||||
|
||||
# Resolve movie ID
|
||||
movie_id = args.movie_id
|
||||
if movie_id is None:
|
||||
print(f"Resolving IMDB ID {args.imdb_id} → TMDB…", file=sys.stderr)
|
||||
movie_id = tmdb_id_from_imdb(args.imdb_id, args.tmdb_key)
|
||||
print(f"TMDB movie ID: {movie_id}", file=sys.stderr)
|
||||
|
||||
# Build gallery
|
||||
gallery = build_gallery(
|
||||
movie_id = movie_id,
|
||||
key = args.tmdb_key,
|
||||
embed_bin = embed_bin,
|
||||
detector = detector,
|
||||
arcface = arcface,
|
||||
max_actors = args.max_actors,
|
||||
images_per_actor = args.images_per_actor,
|
||||
image_root = image_root,
|
||||
)
|
||||
|
||||
n_actors = len(gallery["actors"])
|
||||
n_embeddings = sum(len(a["embeddings"]) for a in gallery["actors"])
|
||||
print(f"\nGallery: {n_actors} actors, {n_embeddings} total embeddings",
|
||||
file=sys.stderr)
|
||||
|
||||
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 __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user