131 lines
4.7 KiB
Python
131 lines
4.7 KiB
Python
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))
|