dtourolle
65fee74585
movienet_eval: replace the per-element dot() with numpy — actor references are loaded once as an ndarray and scored with a single matmul, keeping a whole-library gallery fast. movienet_prep: count and report frames referenced by annotations but absent from Image.zip instead of skipping them silently.
127 lines
4.2 KiB
Python
127 lines
4.2 KiB
Python
#!/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 \
|
|
[--build-dir build]
|
|
|
|
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 sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from sae_embed_loader import load_embedder
|
|
|
|
|
|
def load_gallery(path: str) -> dict[str, dict]:
|
|
"""Return {imdb_id: {"name": str, "refs": np.ndarray[n_refs, dim]}}."""
|
|
with open(path) as f:
|
|
data = json.load(f)
|
|
return {a["imdb_id"]: {"name": a["name"],
|
|
"refs": np.asarray(a["embeddings"], dtype=np.float32)}
|
|
for a in data["actors"]}
|
|
|
|
|
|
def match(embedding: list[float], gallery: dict[str, dict]) -> tuple[str, float, dict[str, float]]:
|
|
"""Return (best_imdb_id, best_similarity, {imdb_id: similarity}).
|
|
|
|
Similarity to an actor is the max dot product over that actor's reference
|
|
embeddings; vectorised with numpy so a whole-library gallery stays fast.
|
|
"""
|
|
vec = np.asarray(embedding, dtype=np.float32)
|
|
scores: dict[str, float] = {
|
|
imdb_id: float((actor["refs"] @ vec).max())
|
|
for imdb_id, actor in gallery.items()
|
|
}
|
|
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("--build-dir", default="build",
|
|
help="Build directory containing the sae_embed module (default: build)")
|
|
p.add_argument("--models-dir", default="models",
|
|
help="Directory containing ONNX models (default: models/)")
|
|
args = p.parse_args()
|
|
|
|
embedder = load_embedder(args.build_dir, args.models_dir, args.arcface)
|
|
|
|
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)
|
|
|
|
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)
|
|
|
|
embed_results = [embedder.embed(str(p)) if p.exists() else None for p in crop_paths]
|
|
|
|
predictions = []
|
|
n_det_fail = 0
|
|
n_correct = 0
|
|
|
|
for entry, result in zip(gt_entries, embed_results):
|
|
detection_failed = result is None or not result.ok
|
|
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()
|