Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/data/alignment.py

@@ -0,0 +1,114 @@
+"""Time-based alignment utilities for audio, video frames, and frame labels.
+
+The model runs on a fixed 25 Hz timeline. Source videos remain at their native
+frame rate; frames and labels are sampled onto the model timeline at load time.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+@dataclass(frozen=True)
+class AlignmentConfig:
+    model_fps: float = 25.0
+    audio_sample_rate: int = 16_000
+
+    def __post_init__(self) -> None:
+        if self.model_fps <= 0:
+            raise ValueError("model_fps must be positive")
+        if self.audio_sample_rate <= 0:
+            raise ValueError("audio_sample_rate must be positive")
+
+
+@dataclass(frozen=True)
+class AVSample:
+    sample_id: str
+    audio_path: Path
+    visual_path: Path
+    source_fps: float
+    labels: Sequence[int]
+    start_sec: float = 0.0
+
+    def __post_init__(self) -> None:
+        if self.source_fps <= 0:
+            raise ValueError(f"{self.sample_id}: source_fps must be positive")
+        if self.start_sec < 0:
+            raise ValueError(f"{self.sample_id}: start_sec cannot be negative")
+
+
+def model_frame_count(source_frames: int, source_fps: float, model_fps: float) -> int:
+    """Return the number of complete model intervals in a source-frame span."""
+    if source_frames < 0 or source_fps <= 0 or model_fps <= 0:
+        raise ValueError("frame counts and frame rates must be valid")
+    return int(np.floor((source_frames / source_fps) * model_fps + 1e-9))
+
+
+def source_indices(
+    source_frames: int, source_fps: float, model_fps: float
+) -> np.ndarray:
+    """Map model-frame centres to the source intervals that contain them."""
+    count = model_frame_count(source_frames, source_fps, model_fps)
+    model_centres = (np.arange(count, dtype=np.float64) + 0.5) / model_fps
+    indices = np.floor(model_centres * source_fps).astype(np.int64)
+    return np.clip(indices, 0, max(source_frames - 1, 0))
+
+
+def model_indices(
+    source_frames: int, source_fps: float, model_fps: float
+) -> np.ndarray:
+    """Map source-frame centres to the model intervals that contain them.
+
+    The inverse direction of `source_indices`, used to lift model-rate
+    predictions back onto the native annotation timeline for evaluation.
+    Always returns exactly `source_frames` indices, so a caller can score one
+    prediction per native frame regardless of the source frame rate.
+    """
+    if source_frames < 0:
+        raise ValueError("source_frames cannot be negative")
+    count = model_frame_count(source_frames, source_fps, model_fps)
+    source_centres = (np.arange(source_frames, dtype=np.float64) + 0.5) / source_fps
+    indices = np.floor(source_centres * model_fps).astype(np.int64)
+    return np.clip(indices, 0, max(count - 1, 0))
+
+
+def align_visual_and_labels(
+    visual: np.ndarray,
+    labels: Sequence[int],
+    source_fps: float,
+    model_fps: float,
+    strict: bool = True,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Sample faces and their labels with one shared timestamp mapping.
+
+    A length disagreement means the face track and the annotation disagree
+    about how long the clip is, which silently shifts every later label.
+    `strict` reports that instead of trimming to the shorter of the two.
+    """
+    labels_array = np.asarray(labels, dtype=np.float32)
+    if strict and len(visual) != len(labels_array):
+        raise ValueError(
+            f"visual has {len(visual)} frames but labels have {len(labels_array)}"
+        )
+    usable_frames = min(len(visual), len(labels_array))
+    if usable_frames == 0:
+        raise ValueError("visual and labels must contain at least one frame")
+
+    indices = source_indices(usable_frames, source_fps, model_fps)
+    return visual[:usable_frames][indices], labels_array[:usable_frames][indices]
+
+
+def fit_audio_to_frames(
+    waveform: torch.Tensor,
+    frame_count: int,
+    sample_rate: int,
+    model_fps: float,
+) -> torch.Tensor:
+    """Pad/crop audio to the exact duration represented by model frames."""
+    target_samples = round(frame_count * sample_rate / model_fps)
+    waveform = waveform[..., :target_samples]
+    return F.pad(waveform, (0, max(0, target_samples - waveform.shape[-1])))