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:
@@ -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()
|
||||
Reference in New Issue
Block a user