Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/data/augment.py

@@ -0,0 +1,105 @@
+"""Train-time augmentation for packed ASD samples.
+
+Every transform here is label-preserving except `substitute_audio`, which
+rewrites the activity labels to silence on purpose. Visual parameters are drawn
+once per chunk and applied to every frame, because the model is temporal and a
+per-frame jitter would destroy the motion it reads.
+"""
+
+import math
+import random
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+
+
+VISUAL_MODES = ("orig", "flip", "crop", "rotate")
+
+
+@dataclass(frozen=True)
+class AugmentConfig:
+    """How aggressively to augment; all probabilities are per chunk."""
+
+    audio_mix_prob: float = 0.5
+    audio_snr_db: tuple[float, float] = (-5.0, 5.0)
+    visual: bool = True
+    crop_scale: tuple[float, float] = (0.7, 1.0)
+    rotate_degrees: float = 15.0
+    # Ablation only: replace the audio outright and call the sample silent.
+    # Stronger than mixing for teaching rejection of off-screen speech, but it
+    # also lets the model key on audio plausibility instead of AV synchrony.
+    substitute_audio_prob: float = 0.0
+
+    @classmethod
+    def from_cfg(cls, cfg: dict | None) -> "AugmentConfig | None":
+        if not cfg:
+            return None
+        snr = cfg.get("audio_snr_db", (-5.0, 5.0))
+        return cls(
+            audio_mix_prob=float(cfg.get("audio_mix_prob", 0.5)),
+            audio_snr_db=(float(snr[0]), float(snr[1])),
+            visual=bool(cfg.get("visual", True)),
+            crop_scale=tuple(cfg.get("crop_scale", (0.7, 1.0))),
+            rotate_degrees=float(cfg.get("rotate_degrees", 15.0)),
+            substitute_audio_prob=float(cfg.get("substitute_audio_prob", 0.0)),
+        )
+
+
+def _power_db(samples: np.ndarray) -> float:
+    return 10.0 * math.log10(float(np.mean(samples.astype(np.float64) ** 2)) + 1e-4)
+
+
+def fit_length(donor: np.ndarray, samples: int) -> np.ndarray:
+    """Wrap or truncate donor audio to the target length."""
+    if len(donor) == 0:
+        return np.zeros(samples, dtype=np.float32)
+    if len(donor) < samples:
+        donor = np.tile(donor, math.ceil(samples / len(donor)))
+    return donor[:samples]
+
+
+def mix_audio(
+    clean: np.ndarray, donor: np.ndarray, snr_db: float
+) -> np.ndarray:
+    """Add donor audio at a target signal-to-noise ratio.
+
+    The label is deliberately unchanged. A negative face with foreign speech
+    added must stay negative, which is what teaches the model to reject speech
+    it cannot see; a positive face stays positive when a competing voice
+    arrives.
+    """
+    donor = fit_length(donor, len(clean))
+    scale = 10.0 ** ((_power_db(clean) - _power_db(donor) - snr_db) / 20.0)
+    return np.clip(clean + donor * scale, -1.0, 1.0)
+
+
+def augment_visual(
+    faces: np.ndarray, config: AugmentConfig, rng: random.Random
+) -> np.ndarray:
+    """Apply one geometric transform to every frame of a chunk."""
+    mode = rng.choice(VISUAL_MODES)
+    if mode == "orig":
+        return faces
+    size = faces.shape[-1]
+    if mode == "flip":
+        return np.ascontiguousarray(faces[:, :, ::-1])
+    if mode == "crop":
+        low, high = config.crop_scale
+        side = max(8, int(size * rng.uniform(low, high)))
+        if side >= size:
+            return faces
+        top = rng.randint(0, size - side)
+        left = rng.randint(0, size - side)
+        window = faces[:, top : top + side, left : left + side]
+        return np.stack(
+            [
+                cv2.resize(frame, (size, size), interpolation=cv2.INTER_LINEAR)
+                for frame in window
+            ]
+        )
+    angle = rng.uniform(-config.rotate_degrees, config.rotate_degrees)
+    matrix = cv2.getRotationMatrix2D((size / 2, size / 2), angle, 1.0)
+    return np.stack(
+        [cv2.warpAffine(frame, matrix, (size, size)) for frame in faces]
+    )