Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/tasks/ava_evaluation.py

@@ -0,0 +1,130 @@
+import csv
+import logging
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import torch
+
+from data.alignment import model_frame_count, model_indices
+from tasks.asd_metrics import ava_active_speaker_map
+
+logger = logging.getLogger(__name__)
+
+MODEL_FPS = 25.0
+
+PREDICTION_FIELDS = [
+    "video_id",
+    "frame_timestamp",
+    "entity_box_x1",
+    "entity_box_y1",
+    "entity_box_x2",
+    "entity_box_y2",
+    "label",
+    "entity_id",
+    "score",
+]
+
+
+def _lift_to_native(entity_id, scores, native_frames, source_fps, model_fps):
+    """Expand model-rate scores to one score per native groundtruth row."""
+    expected_frames = model_frame_count(native_frames, source_fps, model_fps)
+    if len(scores) != expected_frames:
+        raise ValueError(
+            f"{entity_id}: {native_frames} groundtruth rows at {source_fps} fps imply "
+            f"{expected_frames} model frames, got {len(scores)} scores"
+        )
+    indices = model_indices(native_frames, source_fps, model_fps)
+    return scores[torch.as_tensor(indices, dtype=torch.long)]
+
+
+def _count_drifting_entities(rows, fps_by_entity, model_fps):
+    """Count entities whose groundtruth timestamps are not uniformly spaced.
+
+    The lift assumes native frame `j` sits at `t0 + j / source_fps`, matching
+    the positional mapping used when the sample was packed. A few AVA entities
+    have gaps in their annotation timeline; they are scored positionally
+    anyway, but the count belongs in the log so the number stays auditable.
+    """
+    timestamps = defaultdict(list)
+    for row in rows:
+        timestamps[row["entity_id"]].append(float(row["frame_timestamp"]))
+    tolerance = 1.0 / model_fps
+    drifting = 0
+    for entity_id, values in timestamps.items():
+        source_fps = fps_by_entity.get(entity_id)
+        if not source_fps:
+            continue
+        origin = values[0]
+        if any(
+            abs((value - origin) - index / source_fps) > tolerance
+            for index, value in enumerate(values)
+        ):
+            drifting += 1
+    return drifting
+
+
+def export_ava_predictions(
+    groundtruth_path,
+    output_path,
+    scores_by_entity,
+    fps_by_entity=None,
+    model_fps=MODEL_FPS,
+):
+    """Write official AVA predictions and return the mAP over every row.
+
+    `scores_by_entity` holds one score per model frame. When `fps_by_entity` is
+    given, scores are lifted onto each entity's native annotation timeline
+    first, so the exported file has exactly one row per groundtruth row.
+    """
+    groundtruth_path = Path(groundtruth_path)
+    output_path = Path(output_path)
+    rows = list(csv.DictReader(groundtruth_path.open()))
+    expected = Counter(row["entity_id"] for row in rows)
+    missing = expected.keys() - scores_by_entity.keys()
+    extra = scores_by_entity.keys() - expected.keys()
+    if missing or extra:
+        raise ValueError(f"AVA entity mismatch: {len(missing)} missing, {len(extra)} extra")
+
+    if fps_by_entity is not None:
+        absent = expected.keys() - fps_by_entity.keys()
+        if absent:
+            raise ValueError(f"AVA source fps missing for {len(absent)} entities")
+        scores_by_entity = {
+            entity_id: _lift_to_native(
+                entity_id,
+                scores_by_entity[entity_id],
+                expected[entity_id],
+                fps_by_entity[entity_id],
+                model_fps,
+            )
+            for entity_id in expected
+        }
+        drifting = _count_drifting_entities(rows, fps_by_entity, model_fps)
+        if drifting:
+            logger.info(
+                "%d of %d AVA entities have non-uniform groundtruth timestamps; "
+                "scored positionally",
+                drifting,
+                len(expected),
+            )
+
+    for entity_id, count in expected.items():
+        if len(scores_by_entity[entity_id]) != count:
+            raise ValueError(f"{entity_id}: expected {count} scores, got {len(scores_by_entity[entity_id])}")
+
+    positions = defaultdict(int)
+    labels = []
+    scores = []
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=PREDICTION_FIELDS)
+        writer.writeheader()
+        for row in rows:
+            entity_id = row["entity_id"]
+            score = float(scores_by_entity[entity_id][positions[entity_id]])
+            positions[entity_id] += 1
+            prediction = {field: row[field] for field in PREDICTION_FIELDS[:-2]}
+            writer.writerow({**prediction, "label": "SPEAKING_AUDIBLE", "entity_id": entity_id, "score": score})
+            labels.append(row["label"] == "SPEAKING_AUDIBLE")
+            scores.append(score)
+    return ava_active_speaker_map(torch.tensor(scores), torch.tensor(labels))