L'électron rare
aedcb0f01b
Extend action-head to 32 joints (body22 + 10 fingertips), 10 SMPL-X expression PCA scalars, and mouth_open distance. FEATURE_DIM 201→302. MIRROR_MAP extended to 32. Dataset, augment, training, publisher, offline extractor all updated.
86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
"""On-the-fly augmentations for j3d windows."""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
# SMPL-X left/right joint mirror map for 32-joint layout.
|
|
# Body joints 0..21 (unchanged), fingertips 22..31 (L 22..26 <-> R 27..31).
|
|
MIRROR_MAP: tuple[int, ...] = (
|
|
# 22 body (unchanged)
|
|
0,
|
|
2, 1,
|
|
3,
|
|
5, 4,
|
|
6,
|
|
8, 7,
|
|
9,
|
|
11, 10,
|
|
12,
|
|
14, 13,
|
|
15,
|
|
17, 16,
|
|
19, 18,
|
|
21, 20,
|
|
# 10 fingertips: L (22..26) <-> R (27..31)
|
|
27, 28, 29, 30, 31, 22, 23, 24, 25, 26,
|
|
)
|
|
assert len(MIRROR_MAP) == 32
|
|
|
|
|
|
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
|