From 24d1e85b7d8ea4b25cffffd61b047d4498283032 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=27=C3=A9lectron=20rare?=
 <108685187+electron-rare@users.noreply.github.com>
Date: Wed, 13 May 2026 21:40:22 +0200
Subject: [PATCH] feat(data-only-viz): action-head augmentations

Implement on-the-fly spatial and temporal augmentations for multi-HMR
j3d windows: mirror_x (left/right joint swap + x-flip), add_noise,
time_stretch (linear resampling), rotate_y. Task 7 of action-head plan.
All 30 tests pass (26 prior + 4 new augment tests).
---
 data_only_viz/tests/test_augment.py | 45 ++++++++++++++++
 data_only_viz/training/augment.py   | 80 +++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+)
 create mode 100644 data_only_viz/tests/test_augment.py
 create mode 100644 data_only_viz/training/augment.py

diff --git a/data_only_viz/tests/test_augment.py b/data_only_viz/tests/test_augment.py
new file mode 100644
index 0000000..c87b358
--- /dev/null
+++ b/data_only_viz/tests/test_augment.py
@@ -0,0 +1,45 @@
+"""Tests for j3d augmentations."""
+from __future__ import annotations
+
+import numpy as np
+
+WINDOW_LEN = 16
+
+
+def _sample_stack(seed: int = 0) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    return rng.normal(size=(WINDOW_LEN, 22, 3)).astype(np.float32)
+
+
+def test_mirror_swap_left_right_joints() -> None:
+    from data_only_viz.training.augment import mirror_x
+    x = _sample_stack(0)
+    y = mirror_x(x)
+    assert np.allclose(y[..., 0], -x[..., 0][:, [
+        0,2,1,3,5,4,6,8,7,9,11,10,12,14,13,15,17,16,19,18,21,20
+    ]], atol=1e-6)
+
+
+def test_noise_within_sigma() -> None:
+    from data_only_viz.training.augment import add_noise
+    rng = np.random.default_rng(0)
+    x = _sample_stack(0)
+    y = add_noise(x, sigma=0.01, rng=rng)
+    diff = y - x
+    assert np.allclose(diff.std(), 0.01, atol=2e-3)
+
+
+def test_time_stretch_keeps_shape() -> None:
+    from data_only_viz.training.augment import time_stretch
+    x = _sample_stack(0)
+    y = time_stretch(x, factor=0.9, rng=None)
+    assert y.shape == x.shape
+
+
+def test_rotate_y_preserves_distances() -> None:
+    from data_only_viz.training.augment import rotate_y
+    x = _sample_stack(0)
+    y = rotate_y(x, angle_rad=0.3)
+    d_x = np.linalg.norm(x[0, 0] - x[0, 1])
+    d_y = np.linalg.norm(y[0, 0] - y[0, 1])
+    assert abs(d_x - d_y) < 1e-5
diff --git a/data_only_viz/training/augment.py b/data_only_viz/training/augment.py
new file mode 100644
index 0000000..c1b7d93
--- /dev/null
+++ b/data_only_viz/training/augment.py
@@ -0,0 +1,80 @@
+"""On-the-fly augmentations for j3d windows."""
+from __future__ import annotations
+
+import numpy as np
+
+# SMPL-X left/right joint mirror map (subset 22 joints used by Multi-HMR).
+MIRROR_MAP: tuple[int, ...] = (
+    0,
+    2, 1,
+    3,
+    5, 4,
+    6,
+    8, 7,
+    9,
+    11, 10,
+    12,
+    14, 13,
+    15,
+    17, 16,
+    19, 18,
+    21, 20,
+)
+
+
+def mirror_x(stack: np.ndarray) -> np.ndarray:
+    """Mirror across the YZ plane: flip x and swap left↔right joints."""
+    out = stack[:, list(MIRROR_MAP), :].copy()
+    out[..., 0] = -out[..., 0]
+    return out
+
+
+def add_noise(stack: np.ndarray, sigma: float, rng: np.random.Generator) -> np.ndarray:
+    noise = rng.normal(scale=sigma, size=stack.shape).astype(np.float32)
+    return (stack + noise).astype(np.float32, copy=False)
+
+
+def time_stretch(stack: np.ndarray, factor: float,
+                 rng: np.random.Generator | None = None) -> np.ndarray:
+    """Resample the time axis with linear interpolation, keep window_len fixed."""
+    T = stack.shape[0]
+    new_T = int(round(T * factor))
+    new_T = max(2, new_T)
+    src = np.linspace(0.0, T - 1, num=new_T)
+    interp = np.empty((new_T, *stack.shape[1:]), dtype=np.float32)
+    lo = np.floor(src).astype(int)
+    hi = np.minimum(lo + 1, T - 1)
+    frac = (src - lo).astype(np.float32)
+    interp = (1 - frac[:, None, None]) * stack[lo] + frac[:, None, None] * stack[hi]
+    if new_T >= T:
+        start = (new_T - T) // 2
+        return interp[start:start + T].astype(np.float32, copy=False)
+    pad_before = (T - new_T) // 2
+    pad_after = T - new_T - pad_before
+    return np.concatenate([
+        np.repeat(interp[:1], pad_before, axis=0),
+        interp,
+        np.repeat(interp[-1:], pad_after, axis=0),
+    ]).astype(np.float32, copy=False)
+
+
+def rotate_y(stack: np.ndarray, angle_rad: float) -> np.ndarray:
+    """Rotate around Y (vertical) axis."""
+    c, s = np.cos(angle_rad), np.sin(angle_rad)
+    R = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=np.float32)
+    return (stack @ R.T).astype(np.float32, copy=False)
+
+
+def random_augment(stack: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    out = stack
+    if rng.random() < 0.5:
+        out = mirror_x(out)
+    if rng.random() < 0.8:
+        out = add_noise(out, sigma=0.01, rng=rng)
+    if rng.random() < 0.5:
+        factor = float(rng.uniform(0.9, 1.1))
+        out = time_stretch(out, factor=factor, rng=rng)
+    if rng.random() < 0.5:
+        angle = float(rng.uniform(-np.deg2rad(15), np.deg2rad(15)))
+        out = rotate_y(out, angle_rad=angle)
+    return out