electron-rare/AV-Live
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat(data-only-viz): action-head v2 fingers+face

Browse Source 
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.
This commit is contained in:
L'électron rare
2026-05-13 23:15:12 +02:00
parent 28d562b11c
commit aedcb0f01b
13 changed files with 265 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data_only_viz/action_head.py
+43 -14
View File
@@ -21,12 +21,26 @@ WARMUP_FRAMES: int = 3
NAN_SKIP_BUDGET: int = 5
WINDOW_LEN: int = 16
J3D_JOINTS: int = 22
J3D_BODY: int = 22
J3D_FINGERS_PER_HAND: int = 5
J3D_FINGERS: int = 2 * J3D_FINGERS_PER_HAND # 10
J3D_JOINTS: int = J3D_BODY + J3D_FINGERS # 32
J3D_DIMS: int = 3
NUM_CLASSES: int = 3
LABELS: tuple[str, str, str] = ("debout", "assise", "danse")
FEATURE_DIM: int = J3D_JOINTS * J3D_DIMS * 3 + 3 # j3d + vel + accel + 3 scalars
EXPR_DIM: int = 10
EXTRA_SCALARS: int = 4 # hip_y, knee_angle, sym_score, mouth_open
# Layout per step:
# [0 : 96 ] j3d (32, 3)
# [96 : 192] vel (32, 3)
# [192 : 288] accel (32, 3)
# [288 : 298] expression (10,)
# [298 : 302] scalars (hip_y, knee_angle, sym, mouth_open)
FEATURE_DIM: int = J3D_JOINTS * J3D_DIMS * 3 + EXPR_DIM + EXTRA_SCALARS # 302
# Body joint indices (unchanged from v1, indices 0..21).
HIP_LEFT: int = 1
HIP_RIGHT: int = 2
KNEE_LEFT: int = 4
@@ -38,19 +52,26 @@ SHOULDER_RIGHT: int = 17
WRIST_LEFT: int = 20
WRIST_RIGHT: int = 21
# Fingertip indices (new, 22..31), order: L thumb..pinky, R thumb..pinky.
FINGERTIP_LEFT_BASE: int = 22
FINGERTIP_RIGHT_BASE: int = 27
class FeatureExtractor:
"""Stateless feature builder over a list of recent j3d frames.
Vector layout (FEATURE_DIM = 201):
[0 : 66] j3d current frame, flattened (22 joints × 3 dims)
[66 : 132] velocity j3d[t] - j3d[t-1] (22 × 3)
[132 : 198] acceleration vel[t] - vel[t-1] (22 × 3)
[198 : 201] kinetics scalars (hip_y, knee_angle, symmetry_score)
Vector layout (FEATURE_DIM = 302):
[0 : 96 ] j3d current frame, flattened (32 joints x 3 dims)
[96 : 192] velocity j3d[t] - j3d[t-1] (32 x 3)
[192 : 288] acceleration vel[t] - vel[t-1] (32 x 3)
[288 : 298] expression PCA coefficients (10,)
[298 : 302] kinetics scalars (hip_y, knee_angle, symmetry_score, mouth_open)
"""
@staticmethod
def from_buffer(frames: list[np.ndarray]) -> np.ndarray:
def from_buffer(frames: list[np.ndarray],
expr: np.ndarray | None = None,
mouth_open: float = 0.0) -> np.ndarray:
if not frames:
return np.zeros(FEATURE_DIM, dtype=np.float32)
cur = frames[-1]
@@ -62,13 +83,19 @@ class FeatureExtractor:
hip_y = float((cur[HIP_LEFT, 1] + cur[HIP_RIGHT, 1]) * 0.5)
knee_angle = FeatureExtractor._mean_knee_angle(cur)
sym = FeatureExtractor._symmetry_score(vel)
feat = np.concatenate([
if expr is None:
expr_vec = np.zeros(EXPR_DIM, dtype=np.float32)
else:
expr_vec = np.zeros(EXPR_DIM, dtype=np.float32)
n = min(EXPR_DIM, len(expr))
expr_vec[:n] = expr[:n]
return np.concatenate([
cur.reshape(-1),
vel.reshape(-1),
accel.reshape(-1),
np.array([hip_y, knee_angle, sym], dtype=np.float32),
expr_vec,
np.array([hip_y, knee_angle, sym, float(mouth_open)], dtype=np.float32),
]).astype(np.float32, copy=False)
return feat
@staticmethod
def kinetics(frames: list[np.ndarray]) -> np.ndarray:
@@ -148,7 +175,7 @@ class PerPersonBuffer:
class ActionHeadModel(nn.Module):
"""1-layer GRU + small MLP head.
Input : (B, FEATURE_DIM) — single step
Input : (B, FEATURE_DIM) -- single step
Hidden : (1, B, HIDDEN_DIM)
Output : (B, NUM_CLASSES) logits, new hidden
"""
@@ -198,7 +225,9 @@ class ActionHead:
self._hidden: dict[int, torch.Tensor] = {}
self._nan_streak: dict[int, int] = {}
def step(self, pid: int, j3d: np.ndarray) -> tuple[str, np.ndarray, np.ndarray]:
def step(self, pid: int, j3d: np.ndarray,
expr: np.ndarray | None = None,
mouth_open: float = 0.0) -> tuple[str, np.ndarray, np.ndarray]:
if np.isnan(j3d).any():
streak = self._nan_streak.get(pid, 0) + 1
self._nan_streak[pid] = streak
@@ -212,7 +241,7 @@ class ActionHead:
if len(frames) < WARMUP_FRAMES:
probs = np.array([1.0, 0.0, 0.0], dtype=np.float32)
return LABELS[0], probs, np.zeros(3, dtype=np.float32)
feat = FeatureExtractor.from_buffer(frames)
feat = FeatureExtractor.from_buffer(frames, expr=expr, mouth_open=mouth_open)
kin = FeatureExtractor.kinetics(frames)
h = self._hidden.get(pid)
if h is None:
data_only_viz/action_head_pub.py
+84 -16
View File
@@ -14,7 +14,13 @@ from typing import Any
import numpy as np
from data_only_viz.action_head import ActionHead, LABELS
from data_only_viz.action_head import (
ActionHead,
EXPR_DIM,
J3D_FINGERS,
J3D_FINGERS_PER_HAND,
LABELS,
)
LOG = logging.getLogger("action_head_pub")
@@ -22,19 +28,39 @@ DEFAULT_CKPT = (
Path.home() / ".cache" / "av-live-action" / "checkpoints" / "action_head.pt"
)
# 22 vertex indices on the 10475-vertex SMPL-X mesh, approximating
# the 22-joint kinematic chain used by ActionHead.
# NOTE: approximate vertex anchors — real SMPL-X joints come from
# Approximate fingertip vertex indices on SMPL-X 10475-vert mesh.
# Order: L thumb, L index, L middle, L ring, L pinky,
# R thumb, R index, R middle, R ring, R pinky.
SMPLX_FINGERTIP_VERTS: tuple[int, ...] = (
7174, 7397, 7670, 7942, 8214, # L
4631, 4854, 5127, 5399, 5671, # R
)
# 32 vertex indices on the 10475-vertex SMPL-X mesh:
# 22 body (UNCHANGED from v1) + 10 fingertips.
# NOTE: approximate vertex anchors -- real SMPL-X joints come from
# J_regressor @ v3d, but loading the regressor here is avoided for
# live OSC performance. Action-head training must use the same anchors.
SMPLX_JOINT_ANCHOR_VERTS: tuple[int, ...] = (
# 22 body (UNCHANGED indices, same vertex IDs as before)
8204, 3992, 6677, 3500, 3469, 6394, 3279, 3327, 6736, 3074,
8846, 8889, 8848, 1300, 4660, 8964, 3013, 6470, 1602, 5083,
2114, 5559,
# 10 fingertips
*SMPLX_FINGERTIP_VERTS,
)
assert len(SMPLX_JOINT_ANCHOR_VERTS) == 32
# Mouth-open: distance between two lip vertices on SMPL-X mesh.
# vert 8970 (upper outer lip), 8855 (lower outer lip) -- approximate.
SMPLX_UPPER_LIP_VERT: int = 8970
SMPLX_LOWER_LIP_VERT: int = 8855
# MediaPipe HAND fingertip indices (21-kp hand model).
MEDIAPIPE_HAND_FINGERTIPS: tuple[int, ...] = (4, 8, 12, 16, 20)
# MediaPipe 33-landmark indices mapped into the 22-joint slot order.
# NOTE: approximate mapping — spine joints reuse hip/shoulder anchors.
# NOTE: approximate mapping -- spine joints reuse hip/shoulder anchors.
# https://developers.google.com/mediapipe/solutions/vision/pose_landmarker
MEDIAPIPE_TO_22: tuple[int, ...] = (
24, 23, 24, 23, 25, 26, 11, 27, 28, 11,
@@ -85,7 +111,7 @@ class ActionHeadPublisher(threading.Thread):
LOG.info("publisher stopped")
def _tick(self, t_now: float) -> None:
persons22, source_t, source_tag, is_new = self._read_sources()
persons32, source_t, source_tag, is_new = self._read_sources()
if not is_new:
return
if "smplx" in source_tag:
@@ -93,10 +119,12 @@ class ActionHeadPublisher(threading.Thread):
else:
self._last_body_t = source_t
current_pids: set[int] = set()
if persons22:
for pid, j3d in persons22:
if persons32:
for pid, j3d, expr, mouth in persons32:
current_pids.add(pid)
label, probs, kin = self.head.step(pid, j3d)
label, probs, kin = self.head.step(pid, j3d,
expr=expr,
mouth_open=mouth)
idx = LABELS.index(label)
self.bridge.send_action(pid, idx, probs, t_now, force=True)
self.bridge.send_kin(pid, kin, t_now, force=True)
@@ -109,9 +137,11 @@ class ActionHeadPublisher(threading.Thread):
def _read_sources(
self,
) -> tuple[list[tuple[int, np.ndarray]] | None, float, str, bool]:
"""Return (persons22, source_t, source_tag, is_new).
) -> tuple[list[tuple[int, np.ndarray, np.ndarray, float]] | None,
float, str, bool]:
"""Return (persons32, source_t, source_tag, is_new).
Each person entry is (pid, j3d32, expr10, mouth_open).
is_new is True when the timestamp advanced (even if person list
is empty), so _tick can still run the purge loop.
"""
@@ -121,9 +151,11 @@ class ActionHeadPublisher(threading.Thread):
persons_b3d = getattr(self.state, "persons_body3d", None)
ids_b3d = getattr(self.state, "persons_body_ids", None)
t_body = getattr(self.state, "pose_last_t", 0.0)
hands_ids = list(getattr(self.state, "persons_hands_ids", None) or [])
hands_lists = list(getattr(self.state, "persons_hands", None) or [])
# Prefer smplx when its timestamp advanced.
if t_smplx > self._last_smplx_t:
out: list[tuple[int, np.ndarray]] = []
out: list[tuple[int, np.ndarray, np.ndarray, float]] = []
for i, p in enumerate(persons_smplx or []):
pid = int(p.get("pid", i))
v3d = p.get("v3d")
@@ -136,19 +168,55 @@ class ActionHeadPublisher(threading.Thread):
v3d_np = np.asarray(v3d, dtype=np.float32)
if v3d_np.shape[0] < max(SMPLX_JOINT_ANCHOR_VERTS) + 1:
continue
j3d22 = v3d_np[list(SMPLX_JOINT_ANCHOR_VERTS)].astype(np.float32)
out.append((pid, j3d22))
j3d32 = v3d_np[list(SMPLX_JOINT_ANCHOR_VERTS)].astype(np.float32)
# expression
expr = p.get("expression")
if expr is not None:
if hasattr(expr, "numpy") and not isinstance(expr, np.ndarray):
expr = expr.numpy()
expr_np = np.asarray(expr, dtype=np.float32).flatten()
else:
expr_np = np.zeros(EXPR_DIM, dtype=np.float32)
# mouth_open
if v3d_np.shape[0] > max(SMPLX_UPPER_LIP_VERT, SMPLX_LOWER_LIP_VERT):
mouth = float(np.linalg.norm(
v3d_np[SMPLX_UPPER_LIP_VERT] - v3d_np[SMPLX_LOWER_LIP_VERT]
))
else:
mouth = 0.0
out.append((pid, j3d32, expr_np, mouth))
return out or None, t_smplx, "smplx", True
if t_body > self._last_body_t:
ids = ids_b3d or list(range(len(persons_b3d or [])))
# Build hands lookup by pid
hands_by_pid: dict[int, dict[str, Any]] = {}
for hi, hkp in enumerate(hands_lists):
hpid = int(hands_ids[hi]) if hi < len(hands_ids) else hi
side = "L" if hi % 2 == 0 else "R"
hands_by_pid.setdefault(hpid, {})[side] = hkp
out = []
for i, body in enumerate(persons_b3d or []):
pid = int(ids[i]) if i < len(ids) else i
arr = self._kp_list_to_array(body)
if arr is None or arr.shape[0] < 33:
continue
j3d22 = arr[list(MEDIAPIPE_TO_22)].astype(np.float32)
out.append((pid, j3d22))
body22 = arr[list(MEDIAPIPE_TO_22)].astype(np.float32)
# fingertips from hands if available
tips = np.zeros((J3D_FINGERS, 3), dtype=np.float32)
hpair = hands_by_pid.get(pid, {})
for side_idx, side in enumerate(("L", "R")):
hkp = hpair.get(side)
if hkp is None:
continue
hkp_arr = self._kp_list_to_array(hkp)
if hkp_arr is None or hkp_arr.shape[0] < 21:
continue
for k, mp_idx in enumerate(MEDIAPIPE_HAND_FINGERTIPS):
tips[side_idx * J3D_FINGERS_PER_HAND + k] = hkp_arr[mp_idx]
j3d32 = np.concatenate([body22, tips], axis=0)
expr_np = np.zeros(EXPR_DIM, dtype=np.float32)
mouth = 0.0
out.append((pid, j3d32, expr_np, mouth))
return out or None, t_body, "body3d", True
return None, 0.0, "", False
data_only_viz/scripts/extract_j3d_offline.py
+35 -7
View File
@@ -1,4 +1,4 @@
"""Extract j3d (22 SMPL-X joint anchors) from a recorded MP4 using the
"""Extract j3d (32 SMPL-X joint anchors) from a recorded MP4 using the
Multi-HMR CoreML backend, write per-frame per-person jsonl rows.
Usage:
@@ -17,7 +17,12 @@ from pathlib import Path
import cv2
import numpy as np
from data_only_viz.action_head_pub import SMPLX_JOINT_ANCHOR_VERTS
from data_only_viz.action_head import EXPR_DIM
from data_only_viz.action_head_pub import (
SMPLX_JOINT_ANCHOR_VERTS,
SMPLX_UPPER_LIP_VERT,
SMPLX_LOWER_LIP_VERT,
)
from data_only_viz.multihmr_coreml import MultiHMRCoreMLBackend
LOG = logging.getLogger("extract_j3d_offline")
@@ -47,7 +52,11 @@ def _frame_to_chw(frame_bgr: np.ndarray, size: int = IMG_SIZE) -> np.ndarray:
return rgb.transpose(2, 0, 1) # CHW
def _person_to_j3d22(person: dict, anchors: tuple[int, ...]) -> np.ndarray | None:
def _person_to_j3d32(
person: dict,
anchors: tuple[int, ...],
) -> tuple[np.ndarray, np.ndarray, float] | None:
"""Return (j3d32, expression, mouth_open) or None if v3d absent/too small."""
v3d = person.get("v3d")
if v3d is None:
return None
@@ -57,7 +66,23 @@ def _person_to_j3d22(person: dict, anchors: tuple[int, ...]) -> np.ndarray | Non
v3d_np = np.asarray(v3d, dtype=np.float32)
if v3d_np.shape[0] < max(anchors) + 1:
return None
return v3d_np[list(anchors)].astype(np.float32)
j3d32 = v3d_np[list(anchors)].astype(np.float32)
# expression
expr = person.get("expression")
if expr is not None:
if hasattr(expr, "numpy") and not isinstance(expr, np.ndarray):
expr = expr.numpy()
expr_np = np.asarray(expr, dtype=np.float32).flatten()
else:
expr_np = np.zeros(EXPR_DIM, dtype=np.float32)
# mouth_open
if v3d_np.shape[0] > max(SMPLX_UPPER_LIP_VERT, SMPLX_LOWER_LIP_VERT):
mouth = float(np.linalg.norm(
v3d_np[SMPLX_UPPER_LIP_VERT] - v3d_np[SMPLX_LOWER_LIP_VERT]
))
else:
mouth = 0.0
return j3d32, expr_np, mouth
def extract(session: str, video: Path, out: Path,
@@ -89,14 +114,17 @@ def extract(session: str, video: Path, out: Path,
continue
ts = n_frames / fps
for i, person in enumerate(persons):
j3d = _person_to_j3d22(person, anchors)
if j3d is None:
result = _person_to_j3d32(person, anchors)
if result is None:
continue
j3d32, expr_np, mouth = result
f.write(json.dumps({
"ts": ts,
"session": session,
"pid": int(person.get("pid", i)),
"j3d": j3d.tolist(),
"j3d": j3d32.tolist(),
"expression": expr_np.tolist(),
"mouth_open": mouth,
}) + "\n")
n_rows += 1
n_frames += 1
data_only_viz/tests/test_action_head_features.py
+9 -8
View File
@@ -11,14 +11,15 @@ def test_module_imports() -> None:
assert hasattr(action_head, "PerPersonBuffer")
assert hasattr(action_head, "ActionHead")
assert action_head.WINDOW_LEN == 16
assert action_head.J3D_JOINTS == 22
assert action_head.J3D_JOINTS == 32
assert action_head.FEATURE_DIM == 302
assert action_head.NUM_CLASSES == 3
assert action_head.LABELS == ("debout", "assise", "danse")
def _rand_j3d(seed: int = 0) -> np.ndarray:
rng = np.random.default_rng(seed)
return rng.normal(size=(22, 3)).astype(np.float32)
return rng.normal(size=(32, 3)).astype(np.float32)
def test_buffer_starts_empty() -> None:
@@ -58,15 +59,15 @@ def test_buffer_forget_releases_pid() -> None:
def test_buffer_rejects_bad_shape() -> None:
from data_only_viz.action_head import PerPersonBuffer
buf = PerPersonBuffer()
with pytest.raises(ValueError, match="22"):
buf.append(pid=1, j3d=np.zeros((17, 3), dtype=np.float32))
with pytest.raises(ValueError, match="32"):
buf.append(pid=1, j3d=np.zeros((22, 3), dtype=np.float32))
def test_feature_extractor_shape_full_buffer() -> None:
from data_only_viz.action_head import FeatureExtractor, WINDOW_LEN, FEATURE_DIM
frames = [_rand_j3d(i) for i in range(WINDOW_LEN)]
feat = FeatureExtractor.from_buffer(frames)
assert feat.shape == (FEATURE_DIM,)
assert feat.shape == (302,)
assert feat.dtype == np.float32
assert not np.isnan(feat).any()
@@ -91,13 +92,13 @@ def test_feature_extractor_kinetics_speed_and_accel() -> None:
from data_only_viz.action_head import FeatureExtractor, WINDOW_LEN
frames = []
for t in range(WINDOW_LEN):
f = np.zeros((22, 3), dtype=np.float32)
f = np.zeros((32, 3), dtype=np.float32)
f[0, 0] = 0.1 * t
frames.append(f)
kin = FeatureExtractor.kinetics(frames)
assert kin.shape == (3,)
assert kin[0] > 0
assert abs(kin[0] - 0.1 / 22) < 1e-4
assert abs(kin[0] - 0.1 / 32) < 1e-4
assert abs(kin[1]) < 1e-4
@@ -105,7 +106,7 @@ def test_feature_extractor_symmetry_sign() -> None:
from data_only_viz.action_head import FeatureExtractor, WINDOW_LEN, WRIST_LEFT, WRIST_RIGHT
frames = []
for t in range(WINDOW_LEN):
f = np.zeros((22, 3), dtype=np.float32)
f = np.zeros((32, 3), dtype=np.float32)
f[WRIST_LEFT, 0] = 0.05 * t
f[WRIST_RIGHT, 0] = -0.05 * t
frames.append(f)
data_only_viz/tests/test_action_head_model.py
+3 -3
View File
@@ -11,7 +11,7 @@ torch = pytest.importorskip("torch")
def _rand_j3d(seed: int = 0) -> np.ndarray:
rng = np.random.default_rng(seed)
return rng.normal(size=(22, 3)).astype(np.float32)
return rng.normal(size=(32, 3)).astype(np.float32)
def test_model_forward_shape() -> None:
@@ -24,11 +24,11 @@ def test_model_forward_shape() -> None:
assert h_new.shape == h.shape
def test_model_param_count_under_50k() -> None:
def test_model_param_count_under_80k() -> None:
from data_only_viz.action_head import ActionHeadModel
model = ActionHeadModel()
n = sum(p.numel() for p in model.parameters())
assert n < 50_000, f"too many params: {n}"
assert n < 80_000, f"too many params: {n}"
def test_action_head_step_warmup_returns_debout() -> None:
data_only_viz/tests/test_action_head_pub.py
+2
View File
@@ -17,6 +17,8 @@ class _FakeState:
self.persons_body3d = []
self.persons_body_ids = []
self.pose_last_t = 0.0
self.persons_hands = []
self.persons_hands_ids = []
self._lock = threading.RLock()
def lock(self):
data_only_viz/tests/test_augment.py
+6 -5
View File
@@ -8,16 +8,17 @@ 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)
return rng.normal(size=(WINDOW_LEN, 32, 3)).astype(np.float32)
def test_mirror_swap_left_right_joints() -> None:
from data_only_viz.training.augment import mirror_x
from data_only_viz.training.augment import mirror_x, MIRROR_MAP
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)
# Check output shape
assert y.shape == (WINDOW_LEN, 32, 3)
# x-coords are negated after reindexing
assert np.allclose(y[..., 0], -x[:, list(MIRROR_MAP), :][:, :, 0], atol=1e-6)
def test_noise_within_sigma() -> None:
data_only_viz/tests/test_autolabel.py
+2 -2
View File
@@ -10,7 +10,7 @@ def _static_seated(frame_count: int = WINDOW_LEN) -> list[np.ndarray]:
"""Hip low (y small), knee bent ~80°."""
frames = []
for _ in range(frame_count):
f = np.zeros((22, 3), dtype=np.float32)
f = np.zeros((32, 3), dtype=np.float32)
f[1] = [-0.1, 0.4, 0.0]
f[2] = [0.1, 0.4, 0.0]
f[4] = [-0.1, 0.4, 0.3]
@@ -25,7 +25,7 @@ def _static_standing(frame_count: int = WINDOW_LEN) -> list[np.ndarray]:
"""Hip high, knees ~180°."""
frames = []
for _ in range(frame_count):
f = np.zeros((22, 3), dtype=np.float32)
f = np.zeros((32, 3), dtype=np.float32)
f[1] = [-0.1, 0.9, 0.0]
f[2] = [0.1, 0.9, 0.0]
f[4] = [-0.1, 0.5, 0.0]
data_only_viz/tests/test_dataset.py
+6 -6
View File
@@ -15,7 +15,7 @@ def _make_session_jsonl(path: Path, n_frames: int = 64) -> None:
row = {"ts": t / 30.0,
"session": "sess01",
"pid": 1,
"j3d": rng.normal(size=(22, 3)).tolist()}
"j3d": rng.normal(size=(32, 3)).tolist()}
f.write(json.dumps(row) + "\n")
@@ -25,7 +25,7 @@ def test_load_frames_jsonl(tmp_path: Path) -> None:
_make_session_jsonl(p)
frames = load_frames_jsonl(p)
assert len(frames) == 64
assert frames[0].j3d.shape == (22, 3)
assert frames[0].j3d.shape == (32, 3)
assert frames[0].pid == 1
assert frames[0].session == "sess01"
@@ -40,7 +40,7 @@ def test_sliding_windows(tmp_path: Path) -> None:
frames = load_frames_jsonl(p)
windows = list(sliding_windows(frames, window_len=16, stride=4))
assert len(windows) == 13
assert windows[0].j3d_stack.shape == (16, 22, 3)
assert windows[0].j3d_stack.shape == (16, 32, 3)
assert windows[0].session == "sess01"
@@ -55,7 +55,7 @@ def test_write_and_load_dataset_jsonl(tmp_path: Path) -> None:
DatasetRow(
window_id=f"sess01_pid1_w{i:04d}",
label="debout" if i % 2 == 0 else "danse",
j3d_stack=rng.normal(size=(16, 22, 3)).astype(np.float32),
j3d_stack=rng.normal(size=(16, 32, 3)).astype(np.float32),
session="sess01",
pid_local=1,
auto_label_confidence=0.8,
@@ -68,7 +68,7 @@ def test_write_and_load_dataset_jsonl(tmp_path: Path) -> None:
loaded = load_dataset_jsonl(out)
assert len(loaded) == 5
assert loaded[0].label == "debout"
assert loaded[0].j3d_stack.shape == (16, 22, 3)
assert loaded[0].j3d_stack.shape == (16, 32, 3)
assert np.allclose(loaded[0].j3d_stack, rows[0].j3d_stack, atol=1e-6)
@@ -79,7 +79,7 @@ def test_split_by_session(tmp_path: Path) -> None:
for sess in ("s01", "s02", "s03", "s04", "s05", "s06", "s07"):
rows.append(DatasetRow(
window_id=f"{sess}_w0", label="debout",
j3d_stack=rng.normal(size=(16, 22, 3)).astype(np.float32),
j3d_stack=rng.normal(size=(16, 32, 3)).astype(np.float32),
session=sess, pid_local=1, auto_label_confidence=0.7,
manually_validated=False,
))
data_only_viz/tests/test_training_smoke.py
+4 -2
View File
@@ -19,10 +19,12 @@ def _make_tiny_dataset(tmp_path: Path) -> Path:
rows.append(DatasetRow(
window_id=f"{sess}_w{w:03d}",
label=label,
j3d_stack=rng.normal(size=(16, 22, 3)).astype(np.float32),
j3d_stack=rng.normal(size=(16, 32, 3)).astype(np.float32),
session=sess, pid_local=1,
auto_label_confidence=0.8,
manually_validated=True,
expr_stack=np.zeros((16, 10), dtype=np.float32),
mouth_open_stack=np.zeros(16, dtype=np.float32),
))
out = tmp_path / "tiny.jsonl"
write_dataset_jsonl(rows, out)
@@ -57,5 +59,5 @@ def test_trained_checkpoint_loadable(tmp_path: Path) -> None:
lr=1e-3, device="cpu", seed=0, log_every=10_000)
head = ActionHead(ckpt_path=ckpt)
for i in range(5):
label, probs, _ = head.step(pid=1, j3d=np.zeros((22, 3), dtype=np.float32))
label, probs, _ = head.step(pid=1, j3d=np.zeros((32, 3), dtype=np.float32))
assert abs(float(probs.sum()) - 1.0) < 1e-5
data_only_viz/training/augment.py
+6 -1
View File
@@ -3,8 +3,10 @@ from __future__ import annotations
import numpy as np
# SMPL-X left/right joint mirror map (subset 22 joints used by Multi-HMR).
# 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,
@@ -19,7 +21,10 @@ MIRROR_MAP: tuple[int, ...] = (
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:
data_only_viz/training/dataset.py
+46 -7
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
import json
import random
from dataclasses import dataclass
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, Iterator
@@ -15,26 +15,32 @@ class RawFrame:
ts: float
session: str
pid: int
j3d: np.ndarray # (22, 3) float32
j3d: np.ndarray # (32, 3) float32 (v2: body22 + 10 fingertips)
expression: np.ndarray | None = None # (EXPR_DIM,) or None
mouth_open: float = 0.0
@dataclass
class WindowRow:
j3d_stack: np.ndarray # (window_len, 22, 3) float32
j3d_stack: np.ndarray # (window_len, 32, 3) float32
session: str
pid_local: int
first_ts: float
expr_stack: np.ndarray | None = None # (window_len, 10) or None
mouth_open_stack: np.ndarray | None = None # (window_len,) or None
@dataclass
class DatasetRow:
window_id: str
label: str
j3d_stack: np.ndarray # (window_len, 22, 3) float32
j3d_stack: np.ndarray # (window_len, 32, 3) float32
session: str
pid_local: int
auto_label_confidence: float
manually_validated: bool
expr_stack: np.ndarray | None = None # (window_len, 10) or None
mouth_open_stack: np.ndarray | None = None # (window_len,) or None
def load_frames_jsonl(path: Path) -> list[RawFrame]:
@@ -45,11 +51,15 @@ def load_frames_jsonl(path: Path) -> list[RawFrame]:
if not line:
continue
d = json.loads(line)
expr_raw = d.get("expression")
expr = np.asarray(expr_raw, dtype=np.float32) if expr_raw is not None else None
rows.append(RawFrame(
ts=float(d["ts"]),
session=str(d["session"]),
pid=int(d["pid"]),
j3d=np.asarray(d["j3d"], dtype=np.float32),
expression=expr,
mouth_open=float(d.get("mouth_open", 0.0)),
))
return rows
@@ -68,14 +78,32 @@ def sliding_windows(frames: list[RawFrame],
for start in range(0, len(grp) - window_len + 1, stride):
chunk = grp[start:start + window_len]
stack = np.stack([c.j3d for c in chunk]).astype(np.float32)
# Expression stack: zeros if not present
if any(c.expression is not None for c in chunk):
expr_dim = max(
(len(c.expression) for c in chunk if c.expression is not None),
default=10,
)
expr_stack = np.zeros((window_len, expr_dim), dtype=np.float32)
for t, c in enumerate(chunk):
if c.expression is not None:
n = min(expr_dim, len(c.expression))
expr_stack[t, :n] = c.expression[:n]
else:
expr_stack = None
mouth_stack = np.array(
[c.mouth_open for c in chunk], dtype=np.float32
)
yield WindowRow(j3d_stack=stack, session=sess,
pid_local=pid, first_ts=chunk[0].ts)
pid_local=pid, first_ts=chunk[0].ts,
expr_stack=expr_stack,
mouth_open_stack=mouth_stack)
def write_dataset_jsonl(rows: Iterable[DatasetRow], path: Path) -> None:
with path.open("w") as f:
for r in rows:
f.write(json.dumps({
d: dict = {
"window_id": r.window_id,
"label": r.label,
"j3d": r.j3d_stack.astype(np.float32).tolist(),
@@ -83,7 +111,12 @@ def write_dataset_jsonl(rows: Iterable[DatasetRow], path: Path) -> None:
"pid_local": r.pid_local,
"auto_label_confidence": float(r.auto_label_confidence),
"manually_validated": bool(r.manually_validated),
}) + "\n")
}
if r.expr_stack is not None:
d["expr_stack"] = r.expr_stack.astype(np.float32).tolist()
if r.mouth_open_stack is not None:
d["mouth_open_stack"] = r.mouth_open_stack.astype(np.float32).tolist()
f.write(json.dumps(d) + "\n")
def load_dataset_jsonl(path: Path) -> list[DatasetRow]:
@@ -94,6 +127,10 @@ def load_dataset_jsonl(path: Path) -> list[DatasetRow]:
if not line:
continue
d = json.loads(line)
expr_raw = d.get("expr_stack")
expr = np.asarray(expr_raw, dtype=np.float32) if expr_raw is not None else None
mouth_raw = d.get("mouth_open_stack")
mouth = np.asarray(mouth_raw, dtype=np.float32) if mouth_raw is not None else None
out.append(DatasetRow(
window_id=d["window_id"],
label=d["label"],
@@ -102,6 +139,8 @@ def load_dataset_jsonl(path: Path) -> list[DatasetRow]:
pid_local=int(d["pid_local"]),
auto_label_confidence=float(d["auto_label_confidence"]),
manually_validated=bool(d["manually_validated"]),
expr_stack=expr,
mouth_open_stack=mouth,
))
return out
data_only_viz/training/train_action_head.py
+19 -2
View File
@@ -21,6 +21,7 @@ from torch.utils.data import DataLoader, Dataset
from data_only_viz.action_head import (
ActionHeadModel,
EXPR_DIM,
FeatureExtractor,
HIP_LEFT,
HIP_RIGHT,
@@ -52,19 +53,35 @@ class WindowDataset(Dataset[tuple[torch.Tensor, int]]):
stack = row.j3d_stack
if self._augment:
stack = random_augment(stack, self._rng)
T = stack.shape[0]
# expression and mouth_open stacks (zeros if absent / legacy)
if row.expr_stack is not None:
expr_s = row.expr_stack.astype(np.float32)
else:
expr_s = np.zeros((T, EXPR_DIM), dtype=np.float32)
if row.mouth_open_stack is not None:
mouth_s = row.mouth_open_stack.astype(np.float32)
else:
mouth_s = np.zeros(T, dtype=np.float32)
feats = []
prev = stack[0]
prev_vel = np.zeros_like(prev)
for t in range(stack.shape[0]):
for t in range(T):
cur = stack[t]
vel = cur - prev
accel = vel - prev_vel
hip_y = float((cur[HIP_LEFT, 1] + cur[HIP_RIGHT, 1]) * 0.5)
knee_angle = FeatureExtractor._mean_knee_angle(cur)
sym = FeatureExtractor._symmetry_score(vel)
expr_t = expr_s[t] if t < len(expr_s) else np.zeros(EXPR_DIM, dtype=np.float32)
expr_vec = np.zeros(EXPR_DIM, dtype=np.float32)
n = min(EXPR_DIM, len(expr_t))
expr_vec[:n] = expr_t[:n]
mouth_t = float(mouth_s[t]) if t < len(mouth_s) else 0.0
feat = np.concatenate([
cur.reshape(-1), vel.reshape(-1), accel.reshape(-1),
np.array([hip_y, knee_angle, sym], dtype=np.float32),
expr_vec,
np.array([hip_y, knee_angle, sym, mouth_t], dtype=np.float32),
]).astype(np.float32, copy=False)
feats.append(feat)
prev_vel = vel