#!/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 \ --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.mat — N annotations per actor Train[i] = [imdb_id (nm...), count, [[img_path, bbox[x,y,w,h], label], ...]] Image//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.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] = {} n_missing_in_zip = 0 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) except KeyError: n_missing_in_zip += 1 # frame referenced by an annotation but absent from Image.zip continue arr = np.frombuffer(data, dtype=np.uint8) img = cv2.imdecode(arr, cv2.IMREAD_COLOR) if img is not None: frame_cache[img_path] = img print(f"[prep] frames loaded: {len(frame_cache)}/{len(needed_paths)}", file=sys.stderr) if n_missing_in_zip: print(f"[prep] frames absent from Image.zip: {n_missing_in_zip}", 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()