106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
"""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]
|
|
)
|