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
+51
@@ -0,0 +1,51 @@
|
||||
#!/usr/bin/env bash
|
||||
# Pre-build TensorRT engines for ArcFace and SCRFD with the same shape profiles
|
||||
# the runtime nodes use. First-run ORT engine builds take 30–90 s per model and
|
||||
# block the pipeline; this script does it offline so cold starts are instant.
|
||||
#
|
||||
# Profiles must match src/arcface_embedder.hpp and src/scrfd_decoder.hpp:
|
||||
# ArcFace : min=1x3x112x112 opt=Nx3x112x112 max=Nx3x112x112 (N = embed batch)
|
||||
# SCRFD : 1x3x640x640 (fixed; we letterbox to this)
|
||||
#
|
||||
# These trtexec-built engines are *not* picked up by the ORT TRT EP cache —
|
||||
# ORT uses its own engine format. The point of this script is:
|
||||
# (a) sanity-check that the ONNX models build under TRT at all;
|
||||
# (b) measure pure inference latency without ORT overhead.
|
||||
# Run scene_analyze normally and ORT will populate ./trt_cache itself.
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
MODELS="$ROOT/models"
|
||||
OUT="$ROOT/trt_cache"
|
||||
mkdir -p "$OUT"
|
||||
|
||||
EMBED_BATCH="${EMBED_BATCH:-4}"
|
||||
ARCFACE_MODEL="${ARCFACE_MODEL:-$MODELS/arcface_w600k_r50.onnx}"
|
||||
SCRFD_MODEL="${SCRFD_MODEL:-$MODELS/scrfd_500m_bnkps.onnx}"
|
||||
|
||||
run() { echo "+ $*"; "$@"; }
|
||||
|
||||
echo "== ArcFace =="
|
||||
run trtexec \
|
||||
--onnx="$ARCFACE_MODEL" \
|
||||
--fp16 \
|
||||
--minShapes=input.1:1x3x112x112 \
|
||||
--optShapes=input.1:${EMBED_BATCH}x3x112x112 \
|
||||
--maxShapes=input.1:${EMBED_BATCH}x3x112x112 \
|
||||
--saveEngine="$OUT/arcface.$(basename "$ARCFACE_MODEL" .onnx).b${EMBED_BATCH}.fp16.engine" \
|
||||
--useCudaGraph
|
||||
|
||||
echo
|
||||
echo "== SCRFD =="
|
||||
run trtexec \
|
||||
--onnx="$SCRFD_MODEL" \
|
||||
--fp16 \
|
||||
--minShapes=input.1:1x3x640x640 \
|
||||
--optShapes=input.1:1x3x640x640 \
|
||||
--maxShapes=input.1:1x3x640x640 \
|
||||
--saveEngine="$OUT/scrfd.$(basename "$SCRFD_MODEL" .onnx).640.fp16.engine" \
|
||||
--useCudaGraph
|
||||
|
||||
echo
|
||||
echo "Engines saved under: $OUT"
|
||||
echo "Look for 'mean: ... ms' in each section for per-call latency."
|
||||
Executable
+71
@@ -0,0 +1,71 @@
|
||||
#!/usr/bin/env bash
|
||||
# Download ONNX models required by scene_analyze and build_gallery.
|
||||
# Run from the project root: bash scripts/download_models.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
MODELS_DIR="${1:-models}"
|
||||
mkdir -p "$MODELS_DIR"
|
||||
|
||||
# ── YuNet face detection ──────────────────────────────────────────────────────
|
||||
YUNET_URL="https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx"
|
||||
YUNET_FILE="$MODELS_DIR/face_detection_yunet_2023mar.onnx"
|
||||
if [ ! -f "$YUNET_FILE" ]; then
|
||||
echo "Downloading YuNet…"
|
||||
curl -L "$YUNET_URL" -o "$YUNET_FILE"
|
||||
else
|
||||
echo "YuNet already present: $YUNET_FILE"
|
||||
fi
|
||||
|
||||
# ── ArcFace face recognition (buffalo_l / w600k_r50) ─────────────────────────
|
||||
# This model is part of InsightFace's buffalo_l pack.
|
||||
# We download and unpack only the recognition model.
|
||||
ARCFACE_FILE="$MODELS_DIR/arcface_w600k_r50.onnx"
|
||||
if [ ! -f "$ARCFACE_FILE" ]; then
|
||||
echo "Downloading ArcFace (buffalo_l)…"
|
||||
TMP_ZIP=$(mktemp /tmp/buffalo_l.XXXXXX.zip)
|
||||
curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip" \
|
||||
-o "$TMP_ZIP"
|
||||
# The zip contains: 1k3d68.onnx 2d106det.onnx det_10g.onnx genderage.onnx w600k_r50.onnx
|
||||
unzip -jo "$TMP_ZIP" "w600k_r50.onnx" -d "$MODELS_DIR"
|
||||
mv "$MODELS_DIR/w600k_r50.onnx" "$ARCFACE_FILE"
|
||||
rm "$TMP_ZIP"
|
||||
else
|
||||
echo "ArcFace already present: $ARCFACE_FILE"
|
||||
fi
|
||||
|
||||
# ── ArcFace face recognition (buffalo_s / w600k_mbf — MobileFaceNet) ─────────
|
||||
# Lighter backbone (13 MB vs 174 MB for R50) — same 512-dim output, faster inference.
|
||||
ARCFACE_MBF_FILE="$MODELS_DIR/arcface_w600k_mbf.onnx"
|
||||
if [ ! -f "$ARCFACE_MBF_FILE" ]; then
|
||||
echo "Downloading ArcFace MobileFaceNet (buffalo_s)…"
|
||||
TMP_ZIP=$(mktemp /tmp/buffalo_s.XXXXXX.zip)
|
||||
curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_s.zip" \
|
||||
-o "$TMP_ZIP"
|
||||
unzip -jo "$TMP_ZIP" "w600k_mbf.onnx" -d "$MODELS_DIR"
|
||||
mv "$MODELS_DIR/w600k_mbf.onnx" "$ARCFACE_MBF_FILE"
|
||||
rm "$TMP_ZIP"
|
||||
else
|
||||
echo "ArcFace MBF already present: $ARCFACE_MBF_FILE"
|
||||
fi
|
||||
|
||||
# ── SCRFD-500MF face detection (InsightFace buffalo_sc) ───────────────────────
|
||||
# buffalo_sc.zip contains det_500m.onnx (SCRFD-500MF with 5 keypoints).
|
||||
# If the unzip fails (file not found in archive), download manually from:
|
||||
# https://huggingface.co/deepinsight/insightface/resolve/main/models/buffalo_sc/det_500m.onnx
|
||||
SCRFD_FILE="$MODELS_DIR/scrfd_500m_bnkps.onnx"
|
||||
if [ ! -f "$SCRFD_FILE" ]; then
|
||||
echo "Downloading SCRFD-500MF (buffalo_sc)…"
|
||||
TMP_ZIP=$(mktemp /tmp/buffalo_sc.XXXXXX.zip)
|
||||
curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_sc.zip" \
|
||||
-o "$TMP_ZIP"
|
||||
unzip -jo "$TMP_ZIP" "det_500m.onnx" -d "$MODELS_DIR"
|
||||
mv "$MODELS_DIR/det_500m.onnx" "$SCRFD_FILE"
|
||||
rm "$TMP_ZIP"
|
||||
else
|
||||
echo "SCRFD-500MF already present: $SCRFD_FILE"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Models ready in $MODELS_DIR/:"
|
||||
ls -lh "$MODELS_DIR"
|
||||
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()
|
||||
@@ -0,0 +1,144 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
movienet_eval.py — embed probe crops and match against a gallery.
|
||||
|
||||
Usage:
|
||||
python scripts/movienet_eval.py \
|
||||
--gallery gallery_r50.json \
|
||||
--arcface models/arcface_w600k_r50.onnx \
|
||||
--gt eval/gt.json \
|
||||
--output eval/predictions_r50.json \
|
||||
[--yunet models/face_detection_yunet_2023mar.onnx] \
|
||||
[--embed-bin build/embed_faces]
|
||||
|
||||
Input (--gt): list of {"crop": <path>, "imdb_id": <str>, "actor_name": <str>}
|
||||
Output: list of {"crop", "gt", "pred", "similarity", "detection_failed", "all_scores"}
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def load_gallery(path: str) -> dict[str, dict]:
|
||||
"""Return {imdb_id: {"name": str, "embeddings": [[float]]}}."""
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
return {a["imdb_id"]: {"name": a["name"], "embeddings": a["embeddings"]}
|
||||
for a in data["actors"]}
|
||||
|
||||
|
||||
def dot(a: list[float], b: list[float]) -> float:
|
||||
return sum(x * y for x, y in zip(a, b))
|
||||
|
||||
|
||||
def embed_images(paths: list[Path], embed_bin: str, yunet: str, arcface: str) -> list[dict | None]:
|
||||
if not paths:
|
||||
return []
|
||||
cmd = [embed_bin, "--yunet", yunet, "--arcface", arcface] + [str(p) for p in 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(paths)
|
||||
try:
|
||||
return json.loads(proc.stdout)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"[error] embed_faces JSON parse error: {e}", file=sys.stderr)
|
||||
return [None] * len(paths)
|
||||
|
||||
|
||||
def match(embedding: list[float], gallery: dict[str, dict]) -> tuple[str, float, dict[str, float]]:
|
||||
"""Return (best_imdb_id, best_similarity, {imdb_id: similarity})."""
|
||||
scores: dict[str, float] = {}
|
||||
for imdb_id, actor in gallery.items():
|
||||
# max similarity across all reference embeddings for this actor
|
||||
scores[imdb_id] = max(dot(embedding, ref) for ref in actor["embeddings"])
|
||||
best_id = max(scores, key=lambda k: scores[k])
|
||||
return best_id, scores[best_id], scores
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--gallery", required=True)
|
||||
p.add_argument("--arcface", required=True)
|
||||
p.add_argument("--gt", required=True)
|
||||
p.add_argument("--output", required=True)
|
||||
p.add_argument("--yunet", default="models/face_detection_yunet_2023mar.onnx")
|
||||
p.add_argument("--embed-bin", default="build/embed_faces")
|
||||
args = p.parse_args()
|
||||
|
||||
gallery = load_gallery(args.gallery)
|
||||
print(f"[eval] gallery: {len(gallery)} actors", file=sys.stderr)
|
||||
|
||||
with open(args.gt) as f:
|
||||
gt_entries = json.load(f)
|
||||
print(f"[eval] probe crops: {len(gt_entries)}", file=sys.stderr)
|
||||
|
||||
# Batch all crops in one embed_faces call to amortise startup cost
|
||||
crop_paths = [Path(e["crop"]) for e in gt_entries]
|
||||
missing = [p for p in crop_paths if not p.exists()]
|
||||
if missing:
|
||||
print(f"[warn] {len(missing)} crop(s) not found on disk, skipping", file=sys.stderr)
|
||||
|
||||
results_raw = embed_images(
|
||||
[p for p in crop_paths if p.exists()],
|
||||
args.embed_bin, args.yunet, args.arcface
|
||||
)
|
||||
|
||||
# Re-index results back to original list (missing files get None)
|
||||
raw_iter = iter(results_raw)
|
||||
embed_results: list[dict | None] = []
|
||||
for p in crop_paths:
|
||||
embed_results.append(next(raw_iter) if p.exists() else None)
|
||||
|
||||
predictions = []
|
||||
n_det_fail = 0
|
||||
n_correct = 0
|
||||
|
||||
for entry, result in zip(gt_entries, embed_results):
|
||||
detection_failed = result is None or result.get("embedding") is None
|
||||
if detection_failed:
|
||||
n_det_fail += 1
|
||||
predictions.append({
|
||||
"crop": entry["crop"],
|
||||
"gt": entry["imdb_id"],
|
||||
"pred": None,
|
||||
"similarity": None,
|
||||
"detection_failed": True,
|
||||
"all_scores": {},
|
||||
})
|
||||
continue
|
||||
|
||||
pred_id, sim, all_scores = match(result["embedding"], gallery)
|
||||
correct = pred_id == entry["imdb_id"]
|
||||
if correct:
|
||||
n_correct += 1
|
||||
|
||||
predictions.append({
|
||||
"crop": entry["crop"],
|
||||
"gt": entry["imdb_id"],
|
||||
"pred": pred_id,
|
||||
"similarity": sim,
|
||||
"detection_failed": False,
|
||||
"all_scores": all_scores,
|
||||
})
|
||||
|
||||
n_total = len(gt_entries)
|
||||
n_evaluated = n_total - n_det_fail
|
||||
rank1 = n_correct / n_evaluated * 100 if n_evaluated else 0
|
||||
print(f"[eval] detection failures: {n_det_fail}/{n_total}", file=sys.stderr)
|
||||
print(f"[eval] rank-1 accuracy: {rank1:.1f}% ({n_correct}/{n_evaluated})", file=sys.stderr)
|
||||
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(out_path, "w") as f:
|
||||
json.dump(predictions, f, indent=2)
|
||||
print(f"[eval] written → {out_path}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,199 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
movienet_prep.py — extract probe crops from MovieNet-PS for actors in our gallery.
|
||||
|
||||
Usage:
|
||||
python scripts/movienet_prep.py \
|
||||
--movienet <movienet_root> \
|
||||
--gallery gallery.json \
|
||||
--output eval/ \
|
||||
[--split Train_app10] \
|
||||
[--margin 0.2] \
|
||||
[--max-per-actor 50]
|
||||
|
||||
MovieNet-PS format (annotation.zip + Image.zip):
|
||||
annotation/test/train_test/Train_app<N>.mat — N annotations per actor
|
||||
Train[i] = [imdb_id (nm...), count, [[img_path, bbox[x,y,w,h], label], ...]]
|
||||
Image/<movie_tt_id>/shot_XXXX_img_Y.jpg — source frames
|
||||
|
||||
Output:
|
||||
eval/probe/ face crops (jpg)
|
||||
eval/gt.json [{crop, imdb_id, actor_name, source_frame}]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import zipfile
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
except ImportError as e:
|
||||
print(f"[error] missing dependency: {e}", file=sys.stderr)
|
||||
print("Install: pip install opencv-python scipy numpy", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# ── MovieNet-PS loader ────────────────────────────────────────────────────────
|
||||
|
||||
def load_movienet_annotations(movienet_root: Path, split: str) -> list[dict]:
|
||||
"""
|
||||
Parse a MovieNet-PS Train_app<N>.mat split.
|
||||
Returns flat list of {"imdb_id", "img_path", "bbox": [x,y,w,h]}.
|
||||
img_path is relative to Image/ inside Image.zip, e.g. tt0047396/shot_0004_img_1.jpg
|
||||
"""
|
||||
mat_path = movienet_root / "annotation" / "test" / "train_test" / f"{split}.mat"
|
||||
if not mat_path.exists():
|
||||
# try extracting from annotation.zip
|
||||
zip_path = movienet_root / "annotation.zip"
|
||||
if not zip_path.exists():
|
||||
raise FileNotFoundError(f"annotation.zip not found in {movienet_root}")
|
||||
inner = f"annotation/test/train_test/{split}.mat"
|
||||
print(f"[prep] extracting {inner} from annotation.zip…", file=sys.stderr)
|
||||
with zipfile.ZipFile(zip_path) as z:
|
||||
z.extract(inner, movienet_root)
|
||||
mat_path = movienet_root / inner
|
||||
|
||||
data = sio.loadmat(str(mat_path))["Train"]
|
||||
annotations = []
|
||||
for row in data:
|
||||
imdb_id = str(row[0].flat[0]) # e.g. "nm0000023"
|
||||
entries = row[2] # array of [path, bbox, label]
|
||||
for entry in entries:
|
||||
img_path = str(entry[0].flat[0]) # e.g. "tt0032138/shot_0003_img_1.jpg"
|
||||
bbox = [float(v) for v in entry[1].flat] # [x, y, w, h]
|
||||
annotations.append({"imdb_id": imdb_id, "img_path": img_path, "bbox": bbox})
|
||||
return annotations
|
||||
|
||||
|
||||
# ── Gallery loader ────────────────────────────────────────────────────────────
|
||||
|
||||
def load_gallery_ids(gallery_path: str) -> dict[str, str]:
|
||||
"""Return {imdb_id: actor_name} for all actors in the gallery."""
|
||||
with open(gallery_path) as f:
|
||||
data = json.load(f)
|
||||
return {a["imdb_id"]: a["name"] for a in data["actors"]}
|
||||
|
||||
|
||||
# ── Crop + save ───────────────────────────────────────────────────────────────
|
||||
|
||||
def crop_face(img: "np.ndarray", bbox: list[float], margin: float) -> "np.ndarray | None":
|
||||
h, w = img.shape[:2]
|
||||
x, y, bw, bh = bbox
|
||||
# expand by margin
|
||||
pad_x = bw * margin
|
||||
pad_y = bh * margin
|
||||
x1 = max(0, int(x - pad_x))
|
||||
y1 = max(0, int(y - pad_y))
|
||||
x2 = min(w, int(x + bw + pad_x))
|
||||
y2 = min(h, int(y + bh + pad_y))
|
||||
crop = img[y1:y2, x1:x2]
|
||||
return crop if crop.size > 0 else None
|
||||
|
||||
|
||||
# ── Main ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--movienet", required=True, help="MovieNet-PS root directory")
|
||||
p.add_argument("--gallery", required=True, help="gallery.json (for actor list)")
|
||||
p.add_argument("--output", default="eval", help="output directory")
|
||||
p.add_argument("--split", default="Train_app10",
|
||||
help="annotation split to use (default: Train_app10)")
|
||||
p.add_argument("--margin", type=float, default=0.2,
|
||||
help="bbox expansion factor (default 0.2 = 20%%)")
|
||||
p.add_argument("--max-per-actor", type=int, default=50,
|
||||
help="cap probe crops per actor (default 50)")
|
||||
args = p.parse_args()
|
||||
|
||||
movienet_root = Path(args.movienet)
|
||||
out_dir = Path(args.output)
|
||||
probe_dir = out_dir / "probe"
|
||||
probe_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
gallery_ids = load_gallery_ids(args.gallery)
|
||||
print(f"[prep] gallery actors: {len(gallery_ids)}", file=sys.stderr)
|
||||
|
||||
annotations = load_movienet_annotations(movienet_root, args.split)
|
||||
print(f"[prep] total annotations in split: {len(annotations)}", file=sys.stderr)
|
||||
|
||||
matched = [a for a in annotations if a["imdb_id"] in gallery_ids]
|
||||
print(f"[prep] annotations matching gallery: {len(matched)}", file=sys.stderr)
|
||||
|
||||
if not matched:
|
||||
print("[error] no overlap between MovieNet and gallery — check IMDb ID format", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Build set of image paths we actually need, then extract from Image.zip in one pass
|
||||
needed_paths = {a["img_path"] for a in matched}
|
||||
image_zip = movienet_root / "Image.zip"
|
||||
frame_cache: dict[str, np.ndarray] = {}
|
||||
|
||||
print(f"[prep] extracting {len(needed_paths)} frames from Image.zip…", file=sys.stderr)
|
||||
with zipfile.ZipFile(image_zip) as zf:
|
||||
for img_path in needed_paths:
|
||||
zip_entry = f"Image/{img_path}"
|
||||
try:
|
||||
data = zf.read(zip_entry)
|
||||
arr = np.frombuffer(data, dtype=np.uint8)
|
||||
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
||||
if img is not None:
|
||||
frame_cache[img_path] = img
|
||||
except KeyError:
|
||||
pass # file missing from zip, skip silently
|
||||
|
||||
print(f"[prep] frames loaded: {len(frame_cache)}/{len(needed_paths)}", file=sys.stderr)
|
||||
|
||||
per_actor_count: dict[str, int] = {}
|
||||
gt_entries = []
|
||||
n_failed = 0
|
||||
|
||||
for ann in matched:
|
||||
imdb_id = ann["imdb_id"]
|
||||
img_path = ann["img_path"]
|
||||
count = per_actor_count.get(imdb_id, 0)
|
||||
if count >= args.max_per_actor:
|
||||
continue
|
||||
|
||||
img = frame_cache.get(img_path)
|
||||
if img is None:
|
||||
n_failed += 1
|
||||
continue
|
||||
|
||||
crop = crop_face(img, ann["bbox"], args.margin)
|
||||
if crop is None:
|
||||
n_failed += 1
|
||||
continue
|
||||
|
||||
crop_name = f"{imdb_id}_{count:04d}.jpg"
|
||||
crop_path = probe_dir / crop_name
|
||||
cv2.imwrite(str(crop_path), crop)
|
||||
|
||||
per_actor_count[imdb_id] = count + 1
|
||||
gt_entries.append({
|
||||
"crop": str(crop_path),
|
||||
"imdb_id": imdb_id,
|
||||
"actor_name": gallery_ids[imdb_id],
|
||||
"source_frame": img_path,
|
||||
})
|
||||
|
||||
gt_path = out_dir / "gt.json"
|
||||
with open(gt_path, "w") as f:
|
||||
json.dump(gt_entries, f, indent=2)
|
||||
|
||||
print(f"[prep] crops saved: {len(gt_entries)}", file=sys.stderr)
|
||||
print(f"[prep] crop failures: {n_failed}", file=sys.stderr)
|
||||
print(f"[prep] actors covered: {len(per_actor_count)}/{len(gallery_ids)}", file=sys.stderr)
|
||||
for imdb_id, name in sorted(gallery_ids.items()):
|
||||
n = per_actor_count.get(imdb_id, 0)
|
||||
status = f"{n} crops" if n else "NO MATCH"
|
||||
print(f" {name:30s} {status}", file=sys.stderr)
|
||||
print(f"[prep] gt.json → {gt_path}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,144 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
movienet_score.py — compare model predictions against ground truth.
|
||||
|
||||
Usage:
|
||||
python scripts/movienet_score.py \
|
||||
--gt eval/gt.json \
|
||||
--predictions eval/predictions_r50.json:R50:167MB \
|
||||
eval/predictions_r18.json:R18:46MB \
|
||||
eval/predictions_mbf.json:MBF:13MB \
|
||||
[--output eval/report.md]
|
||||
|
||||
Each --predictions value is <path>:<label>:<size> (size is display-only).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def load(path: str) -> list[dict]:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def score(predictions: list[dict], gt_map: dict[str, str]) -> dict:
|
||||
n_total = len(predictions)
|
||||
n_det_fail = sum(1 for p in predictions if p["detection_failed"])
|
||||
evaluated = [p for p in predictions if not p["detection_failed"]]
|
||||
|
||||
n_correct = sum(1 for p in evaluated if p["pred"] == p["gt"])
|
||||
rank1 = n_correct / len(evaluated) * 100 if evaluated else 0.0
|
||||
|
||||
sims_correct = [p["similarity"] for p in evaluated if p["pred"] == p["gt"]]
|
||||
mean_sim = sum(sims_correct) / len(sims_correct) if sims_correct else 0.0
|
||||
|
||||
# Per-actor recall
|
||||
per_actor: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0, "name": ""})
|
||||
for p in evaluated:
|
||||
actor_id = p["gt"]
|
||||
per_actor[actor_id]["total"] += 1
|
||||
per_actor[actor_id]["name"] = gt_map.get(actor_id, actor_id)
|
||||
if p["pred"] == actor_id:
|
||||
per_actor[actor_id]["correct"] += 1
|
||||
|
||||
return {
|
||||
"n_total": n_total,
|
||||
"n_det_fail": n_det_fail,
|
||||
"n_evaluated": len(evaluated),
|
||||
"rank1": rank1,
|
||||
"mean_sim_correct": mean_sim,
|
||||
"per_actor": dict(per_actor),
|
||||
}
|
||||
|
||||
|
||||
def render_table(rows: list[dict], headers: list[str]) -> str:
|
||||
col_widths = [max(len(h), max(len(str(r[h])) for r in rows)) for h in headers]
|
||||
sep = "| " + " | ".join("-" * w for w in col_widths) + " |"
|
||||
header = "| " + " | ".join(h.ljust(w) for h, w in zip(headers, col_widths)) + " |"
|
||||
lines = [header, sep]
|
||||
for r in rows:
|
||||
lines.append("| " + " | ".join(str(r[h]).ljust(w) for h, w in zip(headers, col_widths)) + " |")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--gt", required=True)
|
||||
p.add_argument("--predictions", required=True, nargs="+",
|
||||
metavar="PATH:LABEL:SIZE")
|
||||
p.add_argument("--output", default=None)
|
||||
args = p.parse_args()
|
||||
|
||||
gt_entries = load(args.gt)
|
||||
gt_map = {e["imdb_id"]: e["actor_name"] for e in gt_entries}
|
||||
|
||||
models = []
|
||||
for spec in args.predictions:
|
||||
parts = spec.split(":")
|
||||
if len(parts) != 3:
|
||||
print(f"[error] expected PATH:LABEL:SIZE, got: {spec}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
path, label, size = parts
|
||||
preds = load(path)
|
||||
s = score(preds, gt_map)
|
||||
models.append({"label": label, "size": size, "score": s})
|
||||
|
||||
# ── Summary table ────────────────────────────────────────────────────────────
|
||||
summary_rows = []
|
||||
for m in models:
|
||||
s = m["score"]
|
||||
det_fail_pct = s["n_det_fail"] / s["n_total"] * 100 if s["n_total"] else 0
|
||||
summary_rows.append({
|
||||
"Model": m["label"],
|
||||
"Rank-1": f"{s['rank1']:.1f}%",
|
||||
"Det.Fail": f"{det_fail_pct:.1f}%",
|
||||
"Mean-sim": f"{s['mean_sim_correct']:.3f}",
|
||||
"Probes": str(s["n_total"]),
|
||||
"Size": m["size"],
|
||||
})
|
||||
|
||||
summary_table = render_table(
|
||||
summary_rows,
|
||||
["Model", "Rank-1", "Det.Fail", "Mean-sim", "Probes", "Size"]
|
||||
)
|
||||
|
||||
# ── Per-actor table (using first model's actor list as reference) ────────────
|
||||
all_actor_ids = sorted({e["imdb_id"] for e in gt_entries})
|
||||
actor_rows = []
|
||||
for actor_id in all_actor_ids:
|
||||
row = {"Actor": gt_map.get(actor_id, actor_id)}
|
||||
for m in models:
|
||||
pa = m["score"]["per_actor"].get(actor_id, {"correct": 0, "total": 0})
|
||||
recall = pa["correct"] / pa["total"] * 100 if pa["total"] else 0.0
|
||||
row[m["label"]] = f"{recall:.0f}% ({pa['correct']}/{pa['total']})"
|
||||
actor_rows.append(row)
|
||||
|
||||
actor_headers = ["Actor"] + [m["label"] for m in models]
|
||||
actor_table = render_table(actor_rows, actor_headers)
|
||||
|
||||
# ── Assemble report ──────────────────────────────────────────────────────────
|
||||
report = f"""# MovieNet Validation Report
|
||||
|
||||
## Summary
|
||||
|
||||
{summary_table}
|
||||
|
||||
## Per-Actor Recall
|
||||
|
||||
{actor_table}
|
||||
"""
|
||||
|
||||
print(report)
|
||||
if args.output:
|
||||
out = Path(args.output)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
out.write_text(report)
|
||||
print(f"[score] report written → {out}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user