123 lines
4.1 KiB
Python
123 lines
4.1 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
|
|
|
|
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, "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 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("--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()
|