diff --git a/data_only_viz/CLAUDE.md b/data_only_viz/CLAUDE.md index a4434e0..f252139 100644 --- a/data_only_viz/CLAUDE.md +++ b/data_only_viz/CLAUDE.md @@ -35,6 +35,32 @@ Python **3.11+** requis. `pyproject.toml` est la source de vérité — ne jamai - Shaders Metal dans `shaders/` (`.metal`), recompilés au runtime ; topologie mesh (SMPL faces) en binaire dans `mesh_topology.py`. - OSC out : `osc_listener.py` / `pose_bridge.py` — destination `oscope-of` sur `:57123`. +## action-head (classifier action debout/assise/danse) + +Tête de classification d'action streaming au-dessus des j3d SMPL-X (ou body3d MediaPipe en fallback). Implémentée 2026-05-13. + +| Fichier | Rôle | +|---|---| +| `action_head.py` | `ActionHeadModel` (GRU 1L + MLP, 37 811 params, <2 ms/step M5), `ActionHead.step(pid, j3d) → (label, probs, kin)`, `PerPersonBuffer`, `FeatureExtractor` (201-D : j3d + vel + accel + scalaires) | +| `action_head_pub.py` | Publisher thread démarré dans `multi.py` `__init__`. Polle `state.persons_smplx` (préféré) ou `state.persons_body3d` (fallback) à 30 Hz, dédup par timestamp, extrait j3d22 via `SMPLX_JOINT_ANCHOR_VERTS` ou `MEDIAPIPE_TO_22`, émet OSC `/pose/action` + `/pose/kin` + `/pose/enter/leave` | +| `training/{dataset,autolabel,augment,train_action_head,eval,review}.py` | Pipeline complet : jsonl IO + sliding windows + by-session split / règles auto-label + glue CLI / 4 augmentations / training MPS AdamW CE-weighted / confusion matrix + latence micro-bench / TUI textuel pour review manuel | +| `scripts/capture_actions.py` | Webcam → MP4 + timestamps | +| `scripts/extract_j3d_offline.py` | MP4 → jsonl j3d22 via `MultiHMRCoreMLBackend.infer()` directement (pas de refactor worker) | +| `scripts/train_on_studio.sh` | rsync grosmac → bastion electron-server → studio M3 Ultra + uv sync `--extra multihmr` + train MPS + ckpt back | + +Pipeline complet de capture à live : +```bash +uv run python -m data_only_viz.scripts.capture_actions --session sess01 --duration 600 +uv run python -m data_only_viz.scripts.extract_j3d_offline --session sess01 --video ~/.cache/av-live-action/raw/sess01.mp4 +uv run python -m data_only_viz.training.autolabel --frames ~/.cache/av-live-action/raw/sess01.jsonl --out ~/.cache/av-live-action/dataset/auto.jsonl +uv run python -m data_only_viz.training.review --in ~/.cache/av-live-action/dataset/auto.jsonl --out ~/.cache/av-live-action/dataset/dataset.jsonl +./data_only_viz/scripts/train_on_studio.sh --epochs 50 +uv run python -m data_only_viz.training.eval --ckpt ~/.cache/av-live-action/checkpoints/action_head.pt --dataset ~/.cache/av-live-action/dataset/dataset.jsonl +# Live : publisher déjà câblé dans multi.py, aucune action requise +``` + +Checkpoint par défaut : `~/.cache/av-live-action/checkpoints/action_head.pt`. Absent → random init (warmup retourne `debout`). + ## Tests ```bash @@ -43,6 +69,8 @@ uv run pytest tests/ -v Tests TDD-first pour `nlf_worker.py` ; valider avant chaque commit qui touche un worker. +Suite action-head (8 fichiers, 39 tests) : `tests/test_action_head_*.py`, `tests/test_{dataset,autolabel,augment,training_smoke,pose_bridge_action}.py`. Tous doivent rester verts avant chaque commit qui touche `action_head*.py` ou `training/*.py`. + ## Anti-patterns - Ne pas charger un modèle ML sans guard `try/except ImportError` — les optional-extras peuvent manquer. diff --git a/data_only_viz/action_head.py b/data_only_viz/action_head.py new file mode 100644 index 0000000..c455ede --- /dev/null +++ b/data_only_viz/action_head.py @@ -0,0 +1,279 @@ +"""Action classifier head on top of Multi-HMR j3d. + +Streaming GRU-1-layer + MLP per-person, with a 16-frame ring buffer. +Trained windowed (Studio M3 Ultra MPS), inferred streaming (M5 eager CPU). + +Output per step: (label_idx, probs (3,), kin (3,)) where kin is +(speed_m_s, accel_m_s2, symmetry_in_minus1_plus1). +""" +from __future__ import annotations + +from collections import deque +from pathlib import Path + +import numpy as np +import torch +from torch import nn + +HIDDEN_DIM: int = 48 +MLP_HIDDEN: int = 32 +WARMUP_FRAMES: int = 3 +NAN_SKIP_BUDGET: int = 5 + +WINDOW_LEN: int = 16 +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") + +EXPR_DIM: int = 10 +EXTRA_SCALARS: int = 4 # hip_y, knee_angle, sym_score, mouth_open + +# NEW v3 : MediaPipe Hands keypoints block. +HANDS_KP_PER_HAND: int = 21 +HANDS_KP_TOTAL: int = 2 * HANDS_KP_PER_HAND # 42 +HANDS_KP_DIMS: int = 3 +HANDS_KP_FLAT: int = HANDS_KP_TOTAL * HANDS_KP_DIMS # 126 + +# Layout per step (v3) : +# [0 : 96] j3d (32, 3) +# [96 : 192] vel (32, 3) +# [192 : 288] accel (32, 3) +# [288 : 414] hands_kp (42, 3) zero-padded if absent +# [414 : 424] expression (10,) +# [424 : 428] scalars (hip_y, knee_angle, sym, mouth_open) +FEATURE_DIM: int = J3D_JOINTS * J3D_DIMS * 3 + HANDS_KP_FLAT + EXPR_DIM + EXTRA_SCALARS # 428 + +# Body joint indices (unchanged from v1, indices 0..21). +HIP_LEFT: int = 1 +HIP_RIGHT: int = 2 +KNEE_LEFT: int = 4 +KNEE_RIGHT: int = 5 +ANKLE_LEFT: int = 7 +ANKLE_RIGHT: int = 8 +SHOULDER_LEFT: int = 16 +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 = 428, v3): + [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 : 414] hands_kp (42, 3) MediaPipe Hands, zero-padded if absent + [414 : 424] expression PCA coefficients (10,) + [424 : 428] kinetics scalars (hip_y, knee_angle, symmetry_score, mouth_open) + """ + + @staticmethod + def from_buffer(frames: list[np.ndarray], + expr: np.ndarray | None = None, + mouth_open: float = 0.0, + hands_kp: np.ndarray | None = None) -> np.ndarray: + if not frames: + return np.zeros(FEATURE_DIM, dtype=np.float32) + cur = frames[-1] + prev = frames[-2] if len(frames) >= 2 else cur + prev2 = frames[-3] if len(frames) >= 3 else prev + vel = (cur - prev).astype(np.float32, copy=False) + prev_vel = (prev - prev2).astype(np.float32, copy=False) + accel = (vel - prev_vel).astype(np.float32, copy=False) + 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) + # hands block (42, 3) -> 126 + hands_flat = np.zeros(HANDS_KP_FLAT, dtype=np.float32) + if hands_kp is not None: + hk = np.asarray(hands_kp, dtype=np.float32) + if hk.shape == (HANDS_KP_TOTAL, HANDS_KP_DIMS): + hands_flat = hk.reshape(-1).astype(np.float32, copy=False) + # expression + 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), + hands_flat, + expr_vec, + np.array([hip_y, knee_angle, sym, float(mouth_open)], dtype=np.float32), + ]).astype(np.float32, copy=False) + + @staticmethod + def kinetics(frames: list[np.ndarray]) -> np.ndarray: + """Return (speed, accel_mag, symmetry) averaged over the buffer.""" + if len(frames) < 2: + return np.zeros(3, dtype=np.float32) + arr = np.stack(frames).astype(np.float32, copy=False) + diffs = arr[1:] - arr[:-1] + speeds = np.linalg.norm(diffs, axis=-1).mean(axis=-1) + speed = float(speeds.mean()) + if len(frames) >= 3: + ddiffs = diffs[1:] - diffs[:-1] + accel = float(np.linalg.norm(ddiffs, axis=-1).mean()) + else: + accel = 0.0 + sym = FeatureExtractor._symmetry_score(diffs[-1]) + return np.array([speed, accel, sym], dtype=np.float32) + + @staticmethod + def _mean_knee_angle(j3d: np.ndarray) -> float: + """Angle (rad) at left+right knees, averaged.""" + def _angle(hip: int, knee: int, ankle: int) -> float: + v1 = j3d[hip] - j3d[knee] + v2 = j3d[ankle] - j3d[knee] + n1 = np.linalg.norm(v1) + 1e-6 + n2 = np.linalg.norm(v2) + 1e-6 + cos = float(np.dot(v1, v2) / (n1 * n2)) + return float(np.arccos(np.clip(cos, -1.0, 1.0))) + return 0.5 * (_angle(HIP_LEFT, KNEE_LEFT, ANKLE_LEFT) + + _angle(HIP_RIGHT, KNEE_RIGHT, ANKLE_RIGHT)) + + @staticmethod + def _symmetry_score(vel: np.ndarray) -> float: + """Cosine sim between left-arm and mirrored right-arm velocity.""" + left = vel[WRIST_LEFT].copy() + right = vel[WRIST_RIGHT].copy() + right_mirror = right.copy() + right_mirror[0] = -right_mirror[0] + n1 = np.linalg.norm(left) + 1e-6 + n2 = np.linalg.norm(right_mirror) + 1e-6 + return float(np.dot(left, right_mirror) / (n1 * n2)) + + +class PerPersonBuffer: + """Per-pid ring buffer of j3d frames (deque maxlen=WINDOW_LEN).""" + + __slots__ = ("_buffers",) + + def __init__(self) -> None: + self._buffers: dict[int, deque[np.ndarray]] = {} + + def append(self, pid: int, j3d: np.ndarray) -> None: + if j3d.shape != (J3D_JOINTS, J3D_DIMS): + raise ValueError( + f"j3d must be ({J3D_JOINTS}, {J3D_DIMS}), got {j3d.shape}" + ) + dq = self._buffers.get(pid) + if dq is None: + dq = deque(maxlen=WINDOW_LEN) + self._buffers[pid] = dq + dq.append(j3d.astype(np.float32, copy=False)) + + def frames_for(self, pid: int) -> list[np.ndarray]: + dq = self._buffers.get(pid) + return list(dq) if dq is not None else [] + + def forget(self, pid: int) -> None: + self._buffers.pop(pid, None) + + def __len__(self) -> int: + return len(self._buffers) + + def pids(self) -> list[int]: + return list(self._buffers.keys()) + + +class ActionHeadModel(nn.Module): + """1-layer GRU + small MLP head. + + Input : (B, FEATURE_DIM) -- single step + Hidden : (1, B, HIDDEN_DIM) + Output : (B, NUM_CLASSES) logits, new hidden + """ + + def __init__(self) -> None: + super().__init__() + self.gru = nn.GRU(input_size=FEATURE_DIM, + hidden_size=HIDDEN_DIM, + num_layers=1, + batch_first=True) + self.mlp = nn.Sequential( + nn.Linear(HIDDEN_DIM, MLP_HIDDEN), + nn.ReLU(inplace=True), + nn.Linear(MLP_HIDDEN, NUM_CLASSES), + ) + + def init_hidden(self, batch: int = 1, device: str = "cpu") -> torch.Tensor: + return torch.zeros(1, batch, HIDDEN_DIM, device=device) + + def forward(self, x: torch.Tensor, + h: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + out, h_new = self.gru(x.unsqueeze(1), h) + logits = self.mlp(out.squeeze(1)) + return logits, h_new + + +class ActionHead: + """Streaming action classifier per person. + + Use: + head = ActionHead(ckpt_path=...) + label, probs, kin = head.step(pid, j3d) + head.forget(pid) + """ + + def __init__(self, + ckpt_path: Path | None = None, + device: str = "cpu") -> None: + self._device = device + self._model = ActionHeadModel().to(device).eval() + if ckpt_path is not None: + payload = torch.load(ckpt_path, map_location=device, + weights_only=True) + state = payload.get("model_state_dict", payload) + self._model.load_state_dict(state) + self._buffers = PerPersonBuffer() + self._hidden: dict[int, torch.Tensor] = {} + self._nan_streak: dict[int, int] = {} + + def step(self, pid: int, j3d: np.ndarray, + expr: np.ndarray | None = None, + mouth_open: float = 0.0, + hands_kp: np.ndarray | None = None) -> tuple[str, np.ndarray, np.ndarray]: + if np.isnan(j3d).any(): + streak = self._nan_streak.get(pid, 0) + 1 + self._nan_streak[pid] = streak + if streak > NAN_SKIP_BUDGET: + self.forget(pid) + probs = np.array([1.0, 0.0, 0.0], dtype=np.float32) + return LABELS[0], probs, np.zeros(3, dtype=np.float32) + self._nan_streak[pid] = 0 + self._buffers.append(pid, j3d) + frames = self._buffers.frames_for(pid) + 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, expr=expr, mouth_open=mouth_open, + hands_kp=hands_kp) + kin = FeatureExtractor.kinetics(frames) + h = self._hidden.get(pid) + if h is None: + h = self._model.init_hidden(batch=1, device=self._device) + x = torch.from_numpy(feat).unsqueeze(0).to(self._device) + with torch.no_grad(): + logits, h_new = self._model(x, h) + probs_t = torch.softmax(logits, dim=-1).squeeze(0) + self._hidden[pid] = h_new + probs = probs_t.cpu().numpy().astype(np.float32, copy=False) + return LABELS[int(np.argmax(probs))], probs, kin + + def forget(self, pid: int) -> None: + self._buffers.forget(pid) + self._hidden.pop(pid, None) + self._nan_streak.pop(pid, None) diff --git a/data_only_viz/action_head_pub.py b/data_only_viz/action_head_pub.py new file mode 100644 index 0000000..51bd033 --- /dev/null +++ b/data_only_viz/action_head_pub.py @@ -0,0 +1,313 @@ +"""Action-head publisher : reads state.persons_smplx / persons_body3d, +runs ActionHead per pid, emits /pose/action and /pose/kin via pose_bridge. + +Stand-alone thread to avoid touching multi_hmr_worker.py while it +iterates. Polls state at ~30 Hz, deduplicates by smplx_last_t. +""" +from __future__ import annotations + +import logging +import threading +import time +from pathlib import Path +from typing import Any + +import numpy as np + +from data_only_viz.action_head import ( + ActionHead, + EXPR_DIM, + HANDS_KP_DIMS, + HANDS_KP_PER_HAND, + HANDS_KP_TOTAL, + J3D_FINGERS, + J3D_FINGERS_PER_HAND, + LABELS, +) + +LOG = logging.getLogger("action_head_pub") + +DEFAULT_CKPT = ( + Path.home() / ".cache" / "av-live-action" / "checkpoints" / "action_head.pt" +) + +# Canonical SMPL-X fingertip vertex IDs from smplx.vertex_ids.SMPLX_VERTEX_IDS. +# 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, ...] = ( + 5361, 4933, 5058, 5169, 5286, # L : lthumb, lindex, lmiddle, lring, lpinky + 8079, 7669, 7794, 7905, 8022, # R : rthumb, rindex, rmiddle, rring, rpinky +) + +# 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 FaceMesh inner-mouth landmark indices. +# 13 = upper inner mid, 14 = lower inner mid. +MEDIAPIPE_LIP_UPPER_INNER: int = 13 +MEDIAPIPE_LIP_LOWER_INNER: int = 14 + +# 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. +# https://developers.google.com/mediapipe/solutions/vision/pose_landmarker +MEDIAPIPE_TO_22: tuple[int, ...] = ( + 24, 23, 24, 23, 25, 26, 11, 27, 28, 11, + 31, 32, 0, 11, 12, 0, 11, 12, 13, 14, 15, 16, +) + + +class ActionHeadPublisher(threading.Thread): + """Thread that polls state, runs ActionHead per pid, emits OSC.""" + + def __init__(self, state: Any, bridge: Any, + ckpt_path: Path | None = DEFAULT_CKPT, + period_s: float = 1.0 / 30.0) -> None: + super().__init__(daemon=True, name="action-head-pub") + self.state = state + self.bridge = bridge + self.period = period_s + try: + ckpt = ckpt_path if (ckpt_path and ckpt_path.exists()) else None + self.head = ActionHead(ckpt_path=ckpt, device="cpu") + LOG.info("action_head loaded ckpt=%s", + ckpt if ckpt else "") + except Exception as e: + LOG.warning("action_head init failed: %s", e) + self.head = None + self._stop = threading.Event() + self._last_smplx_t = 0.0 + self._last_body_t = 0.0 + self._last_pids: set[int] = set() + + def stop(self) -> None: + self._stop.set() + + def run(self) -> None: + if self.head is None: + LOG.warning("publisher exiting: no action_head") + return + LOG.info("publisher started") + while not self._stop.is_set(): + t0 = time.perf_counter() + try: + self._tick(t0) + except Exception: + LOG.exception("publisher tick failed") + dt = time.perf_counter() - t0 + if dt < self.period: + time.sleep(self.period - dt) + LOG.info("publisher stopped") + + def _tick(self, t_now: float) -> None: + persons32, source_t, source_tag, is_new = self._read_sources() + if not is_new: + return + if "smplx" in source_tag: + self._last_smplx_t = source_t + else: + self._last_body_t = source_t + current_pids: set[int] = set() + if persons32: + for pid, j3d, expr_np, mouth, hands_kp42 in persons32: + current_pids.add(pid) + label, probs, kin = self.head.step(pid, j3d, expr=expr_np, + mouth_open=mouth, + hands_kp=hands_kp42) + 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) + if pid not in self._last_pids: + self.bridge.send_enter(pid=pid) + for gone in self._last_pids - current_pids: + self.head.forget(gone) + self.bridge.send_leave(pid=gone) + self._last_pids = current_pids + + def _read_sources(self) -> tuple[ + list[tuple[int, np.ndarray, np.ndarray, float, np.ndarray]] | None, + float, str, bool, + ]: + """Return (persons32, source_t, source_tag, is_new). + + Each person entry is (pid, j3d32, expr10, mouth_open, hands_kp42x3). + is_new is True when the timestamp advanced (even if person list + is empty), so _tick can still run the purge loop. + """ + with self.state.lock(): + persons_smplx = getattr(self.state, "persons_smplx", None) + t_smplx = getattr(self.state, "smplx_last_t", 0.0) + persons_b3d = getattr(self.state, "persons_body3d", None) + ids_b3d = getattr(self.state, "persons_body_ids", None) + persons_face = getattr(self.state, "persons_face", None) + ids_face = getattr(self.state, "persons_face_ids", None) + persons_hands = getattr(self.state, "persons_hands", None) + ids_hands = getattr(self.state, "persons_hands_ids", None) + t_body = getattr(self.state, "pose_last_t", 0.0) + + # Build pid -> hands_kp(42, 3) map from MediaPipe persons_hands. + hands_by_pid: dict[int, np.ndarray] = self._build_hands_map( + persons_hands or [], ids_hands or [], + ) + # Build pid -> mouth_open scalar from MediaPipe persons_face lips. + face_mouth_by_pid: dict[int, float] = self._build_face_mouth_map( + persons_face or [], ids_face or [], + ) + + # SMPL-X path (preferred) + if t_smplx > self._last_smplx_t: + out: list[tuple[int, np.ndarray, np.ndarray, float, np.ndarray]] = [] + for i, p in enumerate(persons_smplx or []): + pid = int(p.get("pid", i)) + v3d = p.get("v3d") + if v3d is None: + continue + # CoreMLArray wraps a numpy array but has no __array__ + # protocol; unwrap via .numpy() before np.asarray. + if hasattr(v3d, "numpy") and not isinstance(v3d, np.ndarray): + v3d = v3d.numpy() + v3d_np = np.asarray(v3d, dtype=np.float32) + if v3d_np.shape[0] < max(SMPLX_JOINT_ANCHOR_VERTS) + 1: + continue + 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: prefer MediaPipe face lips, fallback SMPL-X v3d. + if pid in face_mouth_by_pid: + mouth = face_mouth_by_pid[pid] + elif 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 + hands_kp42 = hands_by_pid.get( + pid, np.zeros((HANDS_KP_TOTAL, HANDS_KP_DIMS), dtype=np.float32) + ) + out.append((pid, j3d32, expr_np, mouth, hands_kp42)) + return out or None, t_smplx, "smplx", True + + # MediaPipe body3d fallback + if t_body > self._last_body_t: + ids = ids_b3d or list(range(len(persons_b3d or []))) + 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 + body22 = arr[list(MEDIAPIPE_TO_22)].astype(np.float32) + # fingertips from persons_hands if available + tips = np.zeros((J3D_FINGERS, 3), dtype=np.float32) + hands_kp42 = hands_by_pid.get( + pid, np.zeros((HANDS_KP_TOTAL, HANDS_KP_DIMS), dtype=np.float32) + ) + # extract fingertips from hands_kp42 (idx 4,8,12,16,20 each side) + for side_idx in (0, 1): + base = side_idx * HANDS_KP_PER_HAND + for k, mp_tip in enumerate(MEDIAPIPE_HAND_FINGERTIPS): + if base + mp_tip < hands_kp42.shape[0]: + tips[side_idx * J3D_FINGERS_PER_HAND + k] = \ + hands_kp42[base + mp_tip] + j3d32 = np.concatenate([body22, tips], axis=0) + mouth = face_mouth_by_pid.get(pid, 0.0) + expr_np = np.zeros(EXPR_DIM, dtype=np.float32) + out.append((pid, j3d32, expr_np, mouth, hands_kp42)) + return out or None, t_body, "body3d", True + return None, 0.0, "", False + + def _build_hands_map(self, persons_hands: list, + ids_hands: list) -> dict[int, np.ndarray]: + """Combine left+right hand kp arrays per pid into a single (42, 3) array. + + persons_hands is a flat list ; ids_hands maps each hand-list entry to a + pid (and odd/even index indicates which side). When the user's pipeline + keeps a different convention, this helper makes the best effort and + pads zeros for missing sides. + """ + out: dict[int, np.ndarray] = {} + for hi, hkp in enumerate(persons_hands): + if hkp is None: + continue + pid_raw = ids_hands[hi] if hi < len(ids_hands) else hi + try: + pid = int(pid_raw) + except (TypeError, ValueError): + pid = hi + side = hi % 2 # 0 = L, 1 = R + arr = self._kp_list_to_array(hkp) + if arr is None or arr.shape[0] < HANDS_KP_PER_HAND: + continue + slot = out.setdefault( + pid, np.zeros((HANDS_KP_TOTAL, HANDS_KP_DIMS), dtype=np.float32) + ) + base = side * HANDS_KP_PER_HAND + slot[base:base + HANDS_KP_PER_HAND] = arr[:HANDS_KP_PER_HAND] + return out + + def _build_face_mouth_map(self, persons_face: list, + ids_face: list) -> dict[int, float]: + """Compute mouth_open = norm(upper_inner_lip - lower_inner_lip) per pid.""" + out: dict[int, float] = {} + for fi, fkp in enumerate(persons_face): + if fkp is None: + continue + arr = self._kp_list_to_array(fkp) + if arr is None or arr.shape[0] <= MEDIAPIPE_LIP_LOWER_INNER: + continue + upper = arr[MEDIAPIPE_LIP_UPPER_INNER] + lower = arr[MEDIAPIPE_LIP_LOWER_INNER] + mouth = float(np.linalg.norm(upper - lower)) + try: + pid = int(ids_face[fi]) if fi < len(ids_face) else fi + except (TypeError, ValueError): + pid = fi + out[pid] = mouth + return out + + @staticmethod + def _kp_list_to_array(body: Any) -> np.ndarray | None: + """Best-effort conversion of a body keypoint list to (N, 3) array.""" + if body is None: + return None + if isinstance(body, np.ndarray): + return body + try: + return np.asarray( + [ + ( + getattr(kp, "x", kp[0]), + getattr(kp, "y", kp[1]), + getattr(kp, "z", kp[2] if len(kp) > 2 else 0.0), + ) + for kp in body + ], + dtype=np.float32, + ) + except (TypeError, IndexError, AttributeError): + return None diff --git a/data_only_viz/apple_vision_pose.py b/data_only_viz/apple_vision_pose.py index f94041f..2ccb3f5 100644 --- a/data_only_viz/apple_vision_pose.py +++ b/data_only_viz/apple_vision_pose.py @@ -552,32 +552,13 @@ class AppleVisionPoseWorker: if not hasattr(self, "_logged_face_ok_" + region_name): LOG.info("face: region %s count=%d", region_name, count) setattr(self, "_logged_face_ok_" + region_name, True) - # API stable : pointAtIndex_(k) retourne un CGPoint struct. - n = min(count, end - start) - n_written = 0 - for k in range(n): - try: - pt = region.pointAtIndex_(k) - # CGPoint en pyobjc : tuple (x, y) ou struct - try: - nx_bb = float(pt.x); ny_bb = float(pt.y) - except (AttributeError, TypeError): - nx_bb = float(pt[0]); ny_bb = float(pt[1]) - fx = bx + nx_bb * bw - fy_bl = by + ny_bb * bh - kps[start + k] = PoseKp( - x=fx, y=1.0 - fy_bl, z=0.0, c=1.0) - n_written += 1 - except Exception as e: - if not hasattr(self, "_logged_face_pt_err"): - LOG.info("face: pt %s[%d] err: %s (pt=%r)", - region_name, k, e, type(pt).__name__ - if 'pt' in dir() else "??") - self._logged_face_pt_err = True - continue - if n_written > 0 and not hasattr(self, "_logged_face_write_" + region_name): - LOG.info("face: %s wrote %d points", region_name, n_written) - setattr(self, "_logged_face_write_" + region_name, True) + # pyobjc 11 ne sait pas que pointAtIndex_ prend 1 arg, et + # pointsInImageOfSize_ retourne un PyObjCPointer C-array sans + # API d'acces simple. Face parsing depuis Apple Vision est + # actuellement bloque ; on garde MediaPipe (CPU XNNPACK) pour + # face/hand fin tandis que Vision sert body 2D sur ANE. + # Skip pour eviter le spam ObjCPointerWarning a 30 fps. + return # faceContour fill("faceContour", *FACE_OFFSETS["contour"]) @@ -590,13 +571,19 @@ class AppleVisionPoseWorker: fill("nose", *FACE_OFFSETS["nose"]) fill("medianLine", *FACE_OFFSETS["median"]) - # Pupilles : VNFaceLandmarkRegion2D simple (1 point chacune). + # Pupilles : single-point regions ; meme workaround pyobjc. for region_name, idx in (("leftPupil", 81), ("rightPupil", 82)): try: region = getattr(landmarks, region_name)() if region is None or region.pointCount() < 1: continue - pt = region.pointAtIndex_(0) + try: + pts = region.pointsInImageOfSize_((1.0, 1.0)) + except Exception: + pts = region.normalizedPoints() + if not pts: + continue + pt = pts[0] try: px, py = float(pt.x), float(pt.y) except (AttributeError, TypeError): diff --git a/data_only_viz/main.py b/data_only_viz/main.py index a81423b..0be7474 100644 --- a/data_only_viz/main.py +++ b/data_only_viz/main.py @@ -16,6 +16,7 @@ from __future__ import annotations import argparse import logging +import os import signal import sys @@ -265,9 +266,55 @@ class AppDelegate(NSObject): self._opts, "motion_gate", 5.0), camera_index=getattr(self._opts, "camera_index", -1)) self._pose_worker.start() - self._smplx_tcp = SMPLXTCPSender(self._state) + # MESH_RIG=0 disables the 30 fps rigid translation + # rigger from mesh_rigger.py (used to debug deformation + # issues introduced by the hybrid rigging path). + self._smplx_tcp = SMPLXTCPSender( + self._state, + enable_rigging=os.environ.get("MESH_RIG", "1") != "0", + ) self._smplx_tcp.start() LOG.info("worker: Multi-HMR + SMPL-X (mesh dense)") + # Secondary body-pose worker in parallel: AVLiveBody + # gets body keypoints on UDP :57126 alongside the mesh + # on TCP :57130. Default: Apple Vision (ANE-accel, + # body only 19 joints). Set AV_LIVE_PARALLEL_POSE= + # mediapipe to swap to MediaPipe Holistic (CPU + # XNNPACK but provides face + hand + 3D world). + # Defaut: lance BOTH Apple Vision (body 19 joints sur + # ANE, ~30 fps) ET MediaPipe Multi (face 468 + hands 21 + # + pose 3D world sur CPU XNNPACK). Set + # AV_LIVE_PARALLEL_POSE=apple_vision pour ne garder que + # le path ANE (face/hand fin disparait), ou =mediapipe + # pour ne garder que CPU. + parallel = _os.environ.get( + "AV_LIVE_PARALLEL_POSE", "both") + if parallel in ("apple_vision", "both"): + try: + from .apple_vision_pose import AppleVisionPoseWorker + if AppleVisionPoseWorker.is_available(): + self._av_worker = AppleVisionPoseWorker( + self._state, target_fps=30.0, + num_persons=4) + self._av_worker.start() + LOG.info("worker: + Apple Vision body pose " + "(ANE) in parallel") + else: + raise RuntimeError("apple_vision unavailable") + except Exception as e: # noqa: BLE001 + LOG.warning("Apple Vision parallel start failed " + "(%s)", e) + if parallel in ("mediapipe", "both"): + try: + from .multi import MultiWorker + self._mp_worker = MultiWorker( + self._state, num_persons=4) + self._mp_worker.start() + LOG.info("worker: + MediaPipe Multi (3D pose " + "+ face + hand) in parallel") + except Exception as e: # noqa: BLE001 + LOG.warning("MediaPipe parallel start failed " + "(%s)", e) return LOG.info("Multi-HMR indisponible (checkpoints manquants) " "— voir scripts/setup_multihmr.sh") diff --git a/data_only_viz/mesh_rigger.py b/data_only_viz/mesh_rigger.py new file mode 100644 index 0000000..3eda2fc --- /dev/null +++ b/data_only_viz/mesh_rigger.py @@ -0,0 +1,215 @@ +"""Mesh rigging hybride keyframe (Multi-HMR) + delta Apple Vision. + +Multi-HMR produit un mesh SMPL-X dense (10475 verts) tous les ~300 ms +sur M5 (PyTorch MPS ~3.5 fps). Entre deux keyframes, Apple Vision sur +ANE produit 30 fps de body keypoints 2D. On exploite le pelvis 2D de +Vision pour translater rigidement le mesh keyframe et donner une +perception fluide a 30 fps cote launcher RealityKit. + +Limitations connues (premiere iteration) : + - Translation rigide uniquement (pas de rotation, pas de LBS articule) + - Pelvis 2D delta projete en X/Y a profondeur constante (z keyframe) + - Pas de matching d'identite Vision <-> Multi-HMR : on prend la + personne Vision la plus proche du pelvis projete keyframe +""" +from __future__ import annotations + +import math +import threading +import time +from dataclasses import dataclass, field + +import numpy as np + +from .state import PoseKp, SMPLXPerson, State + + +# Indices MediaPipe POSE_LANDMARKS pour les hanches (pelvis 2D = midpoint). +_LEFT_HIP = 23 +_RIGHT_HIP = 24 + +# Focale par defaut Multi-HMR (camera intrinsics typiques utilisees +# dans multi_hmr_worker : focal = IMG_SIZE). +_IMG_SIZE = 672 +_FOCAL = float(_IMG_SIZE) + + +@dataclass +class _Keyframe: + """Snapshot d'un mesh Multi-HMR + reference Vision au moment T.""" + pid: int + t: float + # Mesh world coords (10475, 3) float32 incluant la translation + vertices_3d: np.ndarray + translation: np.ndarray # (3,) world pelvis + vision_pelvis_2d: tuple[float, float] | None # (cx, cy) normalises 0..1 + + +def _pelvis_2d_from_body(body: list[PoseKp]) -> tuple[float, float] | None: + """Midpoint des deux hanches MediaPipe si confidence > 0.""" + if not body or len(body) <= _RIGHT_HIP: + return None + lh, rh = body[_LEFT_HIP], body[_RIGHT_HIP] + if lh.c <= 0.1 or rh.c <= 0.1: + return None + return (0.5 * (lh.x + rh.x), 0.5 * (lh.y + rh.y)) + + +def _vision_pid_match( + keyframe_pelvis_2d: tuple[float, float] | None, + vision_bodies: list[list[PoseKp]], + vision_ids: list[int], +) -> int | None: + """Retourne le pid Vision dont le pelvis 2D est le plus proche du + keyframe pelvis projete. None si rien.""" + if keyframe_pelvis_2d is None or not vision_bodies: + return None + kx, ky = keyframe_pelvis_2d + best_pid: int | None = None + best_d2 = float("inf") + for body, vpid in zip(vision_bodies, vision_ids): + p = _pelvis_2d_from_body(body) + if p is None: + continue + d2 = (p[0] - kx) ** 2 + (p[1] - ky) ** 2 + if d2 < best_d2: + best_d2 = d2 + best_pid = int(vpid) + return best_pid + + +class MeshRigger: + """Rig le mesh SMPL-X keyframe via le delta pelvis Vision. + + Usage : + rigger = MeshRigger(state) + rigged_persons = rigger.apply(state.persons_smplx, + state.persons_body, + t_now) + Thread-safe : ne mute pas le state, retourne une nouvelle liste. + """ + + def __init__(self, state: State, hold_window_s: float = 1.5) -> None: + self.state = state + self.hold_window_s = hold_window_s + self._lock = threading.Lock() + # pid Multi-HMR -> keyframe + self._keyframes: dict[int, _Keyframe] = {} + # pid Multi-HMR -> pid Vision matched (sticky across frames) + self._vision_pid_map: dict[int, int] = {} + + def apply( + self, + persons_smplx: list[SMPLXPerson], + persons_body: list[list[PoseKp]], + persons_body_ids: list[int], + t_now: float, + ) -> list[SMPLXPerson]: + """Retourne une liste SMPLXPerson translatee par delta Vision.""" + # 1) Detect new keyframes (timestamp tracked via state.smplx_last_t) + with self._lock: + current_pids = {p.pid for p in persons_smplx} + # Drop stale keyframes (person disparue) + for old_pid in list(self._keyframes): + if old_pid not in current_pids: + self._keyframes.pop(old_pid, None) + self._vision_pid_map.pop(old_pid, None) + + out: list[SMPLXPerson] = [] + for person in persons_smplx: + kf = self._keyframes.get(person.pid) + # Detect keyframe refresh : translation differs from kf + is_new_kf = (kf is None or not np.allclose( + kf.translation, person.translation, atol=1e-4)) + if is_new_kf: + # Trouver le pid Vision le plus proche pour ce mesh. + # On projette le pelvis world en 2D image-normalized : + # x_img = (X / Z) * focal / IMG_SIZE + 0.5 + pelvis_2d = self._project_pelvis(person.translation) + matched = _vision_pid_match( + pelvis_2d, persons_body, persons_body_ids) + if matched is None: + matched = self._vision_pid_map.get(person.pid) + if matched is not None: + self._vision_pid_map[person.pid] = matched + # Capture du pelvis 2D Vision au moment du keyframe + vp = None + if matched is not None: + try: + i = persons_body_ids.index(matched) + vp = _pelvis_2d_from_body(persons_body[i]) + except (ValueError, IndexError): + vp = None + self._keyframes[person.pid] = _Keyframe( + pid=person.pid, + t=t_now, + vertices_3d=person.vertices_3d.copy(), + translation=person.translation.copy(), + vision_pelvis_2d=vp, + ) + out.append(person) + continue + + # Entre keyframes : applique delta translation depuis + # Vision pelvis 2D actuel vs keyframe pelvis 2D. + if t_now - kf.t > self.hold_window_s: + # Trop ancien, on lache le rig (mesh statique) + out.append(person) + continue + matched_pid = self._vision_pid_map.get(person.pid) + if matched_pid is None or kf.vision_pelvis_2d is None: + out.append(person) + continue + try: + i = persons_body_ids.index(matched_pid) + except ValueError: + out.append(person) + continue + current_vp = _pelvis_2d_from_body(persons_body[i]) + if current_vp is None: + out.append(person) + continue + + # Image-normalized 2D delta -> world XY delta a depth z_kf. + # Pour un pelvis aux coords image (px in [0,1] centre 0.5), + # X_world = (px - 0.5) * IMG_SIZE * Z / focal = (px-0.5)*Z + # (focal=IMG_SIZE). Delta image -> Delta world a Z fixe. + z_kf = float(kf.translation[2]) if abs( + kf.translation[2]) > 1e-3 else 1.0 + dx_img = current_vp[0] - kf.vision_pelvis_2d[0] + dy_img = current_vp[1] - kf.vision_pelvis_2d[1] + dx_world = dx_img * _IMG_SIZE * z_kf / _FOCAL + dy_world = dy_img * _IMG_SIZE * z_kf / _FOCAL + + # Applique a tous les vertices + a translation. + new_verts = kf.vertices_3d.copy() + new_verts[:, 0] += np.float32(dx_world) + new_verts[:, 1] += np.float32(dy_world) + new_transl = kf.translation.copy() + new_transl[0] += np.float32(dx_world) + new_transl[1] += np.float32(dy_world) + + out.append(SMPLXPerson( + pid=person.pid, + vertices_3d=new_verts, + translation=new_transl, + confidence=person.confidence, + betas=person.betas, + expression=person.expression, + )) + return out + + @staticmethod + def _project_pelvis( + translation: np.ndarray, + ) -> tuple[float, float] | None: + """World pelvis (X,Y,Z) -> image-normalized 2D pelvis.""" + z = float(translation[2]) + if abs(z) < 1e-3: + return None + x_img = (float(translation[0]) * _FOCAL / z) / _IMG_SIZE + 0.5 + y_img = (float(translation[1]) * _FOCAL / z) / _IMG_SIZE + 0.5 + # Clamp en [0,1] + if not (0.0 <= x_img <= 1.0 and 0.0 <= y_img <= 1.0): + return None + return (x_img, y_img) diff --git a/data_only_viz/multi_hmr_worker.py b/data_only_viz/multi_hmr_worker.py index 415f851..a00b6dc 100644 --- a/data_only_viz/multi_hmr_worker.py +++ b/data_only_viz/multi_hmr_worker.py @@ -30,6 +30,7 @@ CACHE = Path.home() / ".cache" / "av-live-multihmr" CKPT = CACHE / "checkpoints" / "multiHMR_672_S.pt" SMPLX_PATH = CACHE / "models" / "smplx" / "SMPLX_NEUTRAL.npz" MULTIHMR_REPO = CACHE / "multi-hmr" +COREML_MLPACKAGE = CACHE / "multihmr_full_672_s.mlpackage" IMG_SIZE = 672 N_VERTS = 10475 @@ -41,7 +42,8 @@ class MultiHMRWorker: det_thresh: float = 0.3, nms_kernel_size: int = 5, motion_gate: float = 5.0, - camera_index: int = -1) -> None: + camera_index: int = -1, + backend: str | None = None) -> None: self.state = state self.num_persons = num_persons self.period = 1.0 / max(1.0, target_fps) @@ -55,6 +57,12 @@ class MultiHMRWorker: self.motion_gate = motion_gate # -1 = auto-select Mac BuiltInWideAngleCamera (cf _camera_select) self.camera_index = camera_index + # backend: 'pytorch' (default) or 'coreml'. CoreML uses the + # .mlpackage at COREML_MLPACKAGE, bypasses MPS torch, and runs + # on ANE/GPU/CPU via CoreML.framework natively (3-4x faster). + self.backend = (backend + or os.environ.get("MULTIHMR_BACKEND", "pytorch") + ).strip().lower() self._stop = threading.Event() self._thread: threading.Thread | None = None self._smooth_shape = [ @@ -72,6 +80,9 @@ class MultiHMRWorker: @staticmethod def is_available() -> bool: + backend = os.environ.get("MULTIHMR_BACKEND", "pytorch").strip().lower() + if backend == "coreml": + return COREML_MLPACKAGE.exists() return CKPT.exists() and SMPLX_PATH.exists() and MULTIHMR_REPO.exists() def start(self) -> None: @@ -83,6 +94,12 @@ class MultiHMRWorker: self._stop.set() def _run(self) -> None: + if self.backend == "coreml": + self._run_coreml() + return + self._run_pytorch() + + def _run_pytorch(self) -> None: if str(MULTIHMR_REPO) not in sys.path: sys.path.insert(0, str(MULTIHMR_REPO)) # Multi-HMR demo.py tire pyrender / pyvista (OpenGL offscreen) et @@ -415,3 +432,223 @@ class MultiHMRWorker: cap.stop() LOG.info("multi_hmr worker stopped") + + # ------------------------------------------------------------------ + # CoreML backend + # ------------------------------------------------------------------ + def _run_coreml(self) -> None: + """CoreML inference path (ANE+GPU+CPU via Apple's framework). + + Mirrors _run_pytorch but loads the .mlpackage via pyobjc + the + CoreML.framework, bypassing torch/MPS entirely. ~3-4x faster + on M5 (28.8ms median vs ~100ms with MPS).""" + try: + import cv2 + except ImportError as e: + LOG.error("opencv-python missing: %s", e) + return + try: + from .multihmr_coreml import MultiHMRCoreMLBackend + backend = MultiHMRCoreMLBackend(COREML_MLPACKAGE) + except Exception as e: # noqa: BLE001 + LOG.error("CoreML backend init failed: %s", e) + return + + focal = float(IMG_SIZE) + K_np = np.array([[focal, 0.0, IMG_SIZE / 2.0], + [0.0, focal, IMG_SIZE / 2.0], + [0.0, 0.0, 1.0]], dtype=np.float32) + + from ._av_capture import ( + AVCapture, find_builtin_device, enumerate_devices) + if self.camera_index >= 0: + devs = enumerate_devices() + if self.camera_index >= len(devs): + LOG.error("camera_index %d hors de %d devices", + self.camera_index, len(devs)) + return + info = devs[self.camera_index] + else: + info = find_builtin_device() + if info is None: + LOG.error("aucune BuiltInWideAngleCamera trouvee") + return + cap = AVCapture(info) + if not cap.start(): + LOG.error("AVCapture start failed pour %s", info["name"]) + return + LOG.info("camera ouverte %s (%s) [coreml backend]", + info["name"], info["type"]) + + frame_count = 0 + persons_count = 0 + skipped_static = 0 + next_heartbeat = time.monotonic() + 5.0 + prev_thumb: np.ndarray | None = None + + while not self._stop.is_set(): + t_cap_start = time.monotonic() + ok, frame_bgr = cap.read(timeout_s=0.5) + if not ok or frame_bgr is None: + time.sleep(self.period) + continue + + t_pre_start = time.monotonic() + h, w = frame_bgr.shape[:2] + if (h, w) != (IMG_SIZE, IMG_SIZE): + side = min(h, w) + y0 = (h - side) // 2 + x0 = (w - side) // 2 + frame_bgr = frame_bgr[y0:y0 + side, x0:x0 + side] + frame_bgr = cv2.resize(frame_bgr, (IMG_SIZE, IMG_SIZE)) + + if self.motion_gate > 0: + thumb = cv2.cvtColor( + cv2.resize(frame_bgr, (112, 112)), + cv2.COLOR_BGR2GRAY) + if prev_thumb is not None: + diff_mean = float(np.mean( + cv2.absdiff(thumb, prev_thumb))) + if diff_mean < self.motion_gate: + prev_thumb = thumb + skipped_static += 1 + time.sleep(self.period) + continue + prev_thumb = thumb + + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + img = frame_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0 + + t_inf_start = time.monotonic() + try: + humans = backend.infer(img, K_np, det_thresh=self.det_thresh) + except Exception as e: # noqa: BLE001 + LOG.warning("coreml inference failed: %s", e) + time.sleep(self.period) + continue + + t_post_start = time.monotonic() + t_now = time.monotonic() + frame_count += 1 + persons_count += len(humans) if humans else 0 + if t_now >= next_heartbeat: + fps = frame_count / 5.0 + avg = persons_count / max(1, frame_count) + LOG.info( + "hb[coreml]: %.1f fps, %.2f persons/frame, %d skipped", + fps, avg, skipped_static) + frame_count = 0 + persons_count = 0 + skipped_static = 0 + next_heartbeat = t_now + 5.0 + + if not humans: + with self.state.lock(): + self.state.persons_smplx = [] + time.sleep(self.period) + continue + + # Dedup intra-frame (same logic as pytorch path). + cand: list[tuple[ + float, float, float, float, float, + np.ndarray, int]] = [] + for i, hh in enumerate(humans): + v = hh["v3d"].detach().cpu().numpy() + xmin = float(v[:, 0].min()); ymin = float(v[:, 1].min()) + xmax = float(v[:, 0].max()); ymax = float(v[:, 1].max()) + score = float(hh["scores"].item()) + pelv = hh["transl_pelvis"].detach().cpu().numpy( + ).flatten()[:3] + cand.append((score, xmin, ymin, xmax, ymax, pelv, i)) + cand.sort(key=lambda c: -c[0]) + keep_idx: list[int] = [] + kept: list[tuple[float, float, float, float, np.ndarray]] = [] + for sc, x0, y0, x1, y1, pelv, src_i in cand: + a_area = max(0.0, x1 - x0) * max(0.0, y1 - y0) + drop = False + for (kx0, ky0, kx1, ky1, kpelv) in kept: + ix0 = max(x0, kx0); iy0 = max(y0, ky0) + ix1 = min(x1, kx1); iy1 = min(y1, ky1) + iw = max(0.0, ix1 - ix0); ih = max(0.0, iy1 - iy0) + inter = iw * ih + if a_area <= 0 or inter <= 0: + continue + k_area = (kx1 - kx0) * (ky1 - ky0) + iou = inter / (a_area + k_area - inter + 1e-9) + pelv_d = float(np.linalg.norm(pelv - kpelv)) + if iou > 0.55 and pelv_d < 0.20: + drop = True + break + if not drop: + keep_idx.append(src_i) + kept.append((x0, y0, x1, y1, pelv)) + if len(keep_idx) >= self.num_persons: + break + humans = [humans[i] for i in keep_idx] + n_keep = len(humans) + + bboxes = [] + for hh in humans: + v = hh["v3d"].detach().cpu().numpy() + xmin, ymin = float(v[:, 0].min()), float(v[:, 1].min()) + xmax, ymax = float(v[:, 0].max()), float(v[:, 1].max()) + bboxes.append([PoseKp(x=xmin, y=ymin, c=1.0), + PoseKp(x=xmax, y=ymax, c=1.0)]) + ids = self._tracker.update(bboxes) + + persons: list[SMPLXPerson] = [] + for i, hh in enumerate(humans[:n_keep]): + pid = ids[i] if i < len(ids) else i + if pid < 0: + continue + v3d = hh["v3d"].detach().cpu().numpy() + transl_np = hh["transl_pelvis"].detach().cpu().numpy().flatten() + shape_raw = hh["shape"].detach().cpu().numpy().flatten() + expr_raw = hh["expression"].detach().cpu().numpy().flatten() + + pid_c = pid % self.num_persons + shape_n = min(10, len(shape_raw)) + expr_n = min(10, len(expr_raw)) + shape_smooth = np.zeros(10, dtype=np.float32) + expr_smooth = np.zeros(10, dtype=np.float32) + for k in range(shape_n): + shape_smooth[k] = self._smooth_shape[pid_c][k]( + float(shape_raw[k]), t_now) + for k in range(expr_n): + expr_smooth[k] = self._smooth_expr[pid_c][k]( + float(expr_raw[k]), t_now) + + persons.append(SMPLXPerson( + pid=int(pid), + vertices_3d=np.ascontiguousarray(v3d, dtype=np.float32), + translation=np.ascontiguousarray( + transl_np[:3], dtype=np.float32), + confidence=float(hh["scores"].item()), + betas=np.ascontiguousarray(shape_smooth, dtype=np.float32), + expression=np.ascontiguousarray(expr_smooth, dtype=np.float32), + )) + + with self.state.lock(): + self.state.persons_smplx = persons + self.state.smplx_last_t = t_now + + t_end = time.monotonic() + dt_total = (t_end - t_cap_start) * 1e3 + if LOG.isEnabledFor(logging.DEBUG) or dt_total > 100.0: + LOG.log( + logging.DEBUG if dt_total <= 100.0 else logging.WARNING, + "frame[coreml]: cap=%.1f pre=%.1f inf=%.1f " + "post=%.1fms total=%.1fms", + (t_pre_start - t_cap_start) * 1e3, + (t_inf_start - t_pre_start) * 1e3, + (t_post_start - t_inf_start) * 1e3, + (t_end - t_post_start) * 1e3, + dt_total, + ) + + dt = time.monotonic() - t_cap_start + if dt < self.period: + time.sleep(self.period - dt) + + cap.stop() + LOG.info("multi_hmr coreml worker stopped") diff --git a/data_only_viz/multihmr_coreml.py b/data_only_viz/multihmr_coreml.py new file mode 100644 index 0000000..3992800 --- /dev/null +++ b/data_only_viz/multihmr_coreml.py @@ -0,0 +1,276 @@ +"""Multi-HMR CoreML backend (ANE/GPU/CPU via Apple's CoreML framework). + +Python 3.14 cannot use `coremltools.MLModel` because `libcoremlpython` +and `libmilstoragepython` native extensions are not distributed for +3.14. We load CoreML.framework directly via `objc.loadBundle()` — +same pattern as `coreml_pose.py`. + +Unlike `coreml_pose.py`, this backend does NOT use Vision: Vision is +limited to image inputs and cannot feed a second MLMultiArray (cam_K). +We invoke `MLModel.predictionFromFeatures:error:` directly with a +`MLDictionaryFeatureProvider` wrapping two `MLMultiArray`s. + +Public API: + backend = MultiHMRCoreMLBackend(mlpackage_path) + humans = backend.infer(image_chw_f32, K_33_f32, det_thresh=0.3) + # humans is a list[dict] with the same keys as the PyTorch model + # output. Values are CoreMLArray instances that quack like torch + # tensors (.detach().cpu().numpy() / .item()). +""" +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any + +import numpy as np +import objc +from Foundation import NSURL + +LOG = logging.getLogger("multihmr_coreml") + +DEFAULT_MLPACKAGE = ( + Path.home() / ".cache" / "av-live-multihmr" + / "multihmr_full_672_s.mlpackage" +) + +# Multi-HMR exported with apply_topk(K=4): outputs are fixed shape. +N_PERSONS_FIXED = 4 +N_VERTS = 10475 + +# CoreML output names from the exported .mlpackage. +OUT_V3D = "var_2412" # (4, 10475, 3) +OUT_TRANSL = "var_2415" # (4, 1, 3) +OUT_SCORES = "var_2428" # (4,) +OUT_BETAS = "var_2431" # (4, 10) +OUT_EXPR = "var_2434" # (4, 10) + +# MLMultiArrayDataType raw values (from CoreML headers). +ML_DTYPE_FLOAT32 = 65568 +ML_DTYPE_FLOAT16 = 65552 +ML_DTYPE_DOUBLE = 65600 +ML_DTYPE_INT32 = 131104 + + +_NS: dict[str, Any] = {} +_FRAMEWORKS_LOADED = False + + +def _load_frameworks() -> dict[str, Any]: + global _FRAMEWORKS_LOADED + if _FRAMEWORKS_LOADED: + return _NS + objc.loadBundle("CoreML", _NS, + "/System/Library/Frameworks/CoreML.framework") + _FRAMEWORKS_LOADED = True + return _NS + + +class CoreMLArray: + """Tiny tensor-like adapter so the existing worker hot path can + treat CoreML outputs the same way it treats torch tensors. + + Supports `.detach().cpu().numpy()` and `.item()`. The wrapper is + a no-op around a numpy array; we keep the chain so callers don't + need any conditional branch.""" + + __slots__ = ("_arr",) + + def __init__(self, arr: np.ndarray) -> None: + self._arr = arr + + def detach(self) -> "CoreMLArray": + return self + + def cpu(self) -> "CoreMLArray": + return self + + def numpy(self) -> np.ndarray: + return self._arr + + def item(self) -> float: + return float(self._arr.reshape(-1)[0]) + + @property + def shape(self) -> tuple[int, ...]: + return tuple(self._arr.shape) + + +def _np_to_mlarray(arr: np.ndarray): + """Create a contiguous float32 MLMultiArray from a numpy array. + + We always feed FLOAT32 — even though outputs are FLOAT16, CoreML + will auto-cast on the input side.""" + ns = _load_frameworks() + MLMultiArray = ns["MLMultiArray"] + arr = np.ascontiguousarray(arr, dtype=np.float32) + shape = [int(s) for s in arr.shape] + ml = MLMultiArray.alloc().initWithShape_dataType_error_( + shape, ML_DTYPE_FLOAT32, None) + if ml is None: + raise RuntimeError("MLMultiArray alloc failed") + # Copy bytes through dataPointer (raw void*). pyobjc exposes it as + # a memoryview-like opaque; we use ctypes to memcpy. + import ctypes + ptr = ml.dataPointer() + n_bytes = arr.nbytes + # pyobjc returns either an objc.varlist or a Python int pointer. + addr = int(ptr) if isinstance(ptr, int) else ctypes.cast( + ptr, ctypes.c_void_p).value + if addr is None: + raise RuntimeError("MLMultiArray dataPointer null") + ctypes.memmove(addr, arr.ctypes.data, n_bytes) + return ml + + +def _mlarray_to_np(ml) -> np.ndarray: + """Copy an MLMultiArray (FLOAT16 or FLOAT32) into a numpy float32.""" + import ctypes + shape = tuple(int(s) for s in ml.shape()) + dtype_id = int(ml.dataType()) + count = 1 + for s in shape: + count *= s + ptr = ml.dataPointer() + addr = int(ptr) if isinstance(ptr, int) else ctypes.cast( + ptr, ctypes.c_void_p).value + if addr is None: + raise RuntimeError("MLMultiArray dataPointer null") + if dtype_id == ML_DTYPE_FLOAT16: + raw = (ctypes.c_uint16 * count).from_address(addr) + arr = np.ctypeslib.as_array(raw).view(np.float16).astype(np.float32) + elif dtype_id == ML_DTYPE_FLOAT32: + raw = (ctypes.c_float * count).from_address(addr) + arr = np.ctypeslib.as_array(raw).copy() + elif dtype_id == ML_DTYPE_DOUBLE: + raw = (ctypes.c_double * count).from_address(addr) + arr = np.ctypeslib.as_array(raw).astype(np.float32) + else: + raise RuntimeError(f"unsupported MLMultiArray dtype {dtype_id}") + return arr.reshape(shape) + + +class MultiHMRCoreMLBackend: + """CoreML inference wrapper for Multi-HMR (full_672_s).""" + + def __init__(self, mlpackage_path: Path | None = None) -> None: + self.path = Path(mlpackage_path) if mlpackage_path else DEFAULT_MLPACKAGE + if not self.path.exists(): + raise FileNotFoundError(f"mlpackage missing: {self.path}") + ns = _load_frameworks() + MLModel = ns["MLModel"] + MLModelConfiguration = ns["MLModelConfiguration"] + cfg = MLModelConfiguration.alloc().init() + try: + # MLComputeUnits: 0=CPUOnly, 1=CPUAndGPU, 2=All (ANE+GPU+CPU), + # 3=CPUAndNeuralEngine. Multi-HMR's ANEF compile fails + # (validated 2026-05-13 on M5), and 'All' falls back to a + # slow path (~146ms). CPU+GPU = 28ms = ~35fps on M5. + cfg.setComputeUnits_(1) + except Exception: # noqa: BLE001 + pass + url = NSURL.fileURLWithPath_(str(self.path)) + # .mlpackage must be compiled to .mlmodelc before MLModel can + # load it. compileModelAtURL_error_ returns an NSURL to a + # temp .mlmodelc bundle. + compiled_url = MLModel.compileModelAtURL_error_(url, None) + if compiled_url is None: + raise RuntimeError(f"compileModelAtURL failed for {self.path}") + model = MLModel.modelWithContentsOfURL_configuration_error_( + compiled_url, cfg, None) + if model is None: + raise RuntimeError(f"MLModel load failed for {compiled_url}") + self._model = model + self._ns = ns + LOG.info("Multi-HMR CoreML model loaded (%s, computeUnits=CPU+GPU)", + self.path.name) + + @staticmethod + def is_available(mlpackage_path: Path | None = None) -> bool: + p = Path(mlpackage_path) if mlpackage_path else DEFAULT_MLPACKAGE + if not p.exists(): + return False + try: + _load_frameworks() + return True + except Exception: # noqa: BLE001 + return False + + def _predict(self, image_4d: np.ndarray, K_33: np.ndarray) -> dict: + ns = self._ns + MLDictionaryFeatureProvider = ns["MLDictionaryFeatureProvider"] + MLFeatureValue = ns["MLFeatureValue"] + img_ml = _np_to_mlarray(image_4d) + k_ml = _np_to_mlarray(K_33) + feats = { + "image": MLFeatureValue.featureValueWithMultiArray_(img_ml), + "cam_K": MLFeatureValue.featureValueWithMultiArray_(k_ml), + } + provider = MLDictionaryFeatureProvider.alloc( + ).initWithDictionary_error_(feats, None) + if provider is None: + raise RuntimeError("MLDictionaryFeatureProvider alloc failed") + out = self._model.predictionFromFeatures_error_(provider, None) + if out is None: + raise RuntimeError("MLModel predict failed") + names = [str(n) for n in out.featureNames()] + result = {} + for n in names: + fv = out.featureValueForName_(n) + ml = fv.multiArrayValue() + if ml is None: + continue + result[n] = _mlarray_to_np(ml) + return result + + def infer( + self, + image_chw_float32: np.ndarray, + K_33: np.ndarray, + det_thresh: float = 0.3, + ) -> list[dict]: + """Run a forward pass and return list of humans dicts. + + Args: + image_chw_float32: (3, 672, 672) or (1, 3, 672, 672) in [0,1]. + K_33: (3, 3) or (1, 3, 3) camera intrinsics. + det_thresh: scores threshold; CoreML forwards K=4 always. + + Returns: + list[dict] with keys v3d, transl_pelvis, scores, shape, + expression. Values are CoreMLArray wrappers. + """ + img = np.asarray(image_chw_float32, dtype=np.float32) + if img.ndim == 3: + img = img[np.newaxis, ...] + if img.shape != (1, 3, 672, 672): + raise ValueError(f"image shape {img.shape}, expected (1,3,672,672)") + K = np.asarray(K_33, dtype=np.float32) + if K.ndim == 2: + K = K[np.newaxis, ...] + if K.shape != (1, 3, 3): + raise ValueError(f"K shape {K.shape}, expected (1,3,3)") + + raw = self._predict(img, K) + v3d = raw.get(OUT_V3D) + transl = raw.get(OUT_TRANSL) + scores = raw.get(OUT_SCORES) + betas = raw.get(OUT_BETAS) + expr = raw.get(OUT_EXPR) + if any(x is None for x in (v3d, transl, scores, betas, expr)): + raise RuntimeError( + "missing outputs; got keys=" + ",".join(raw.keys())) + + humans: list[dict] = [] + for k in range(N_PERSONS_FIXED): + sc = float(scores[k]) + if sc < det_thresh: + continue + humans.append({ + "v3d": CoreMLArray(v3d[k]), # (10475, 3) + "transl_pelvis": CoreMLArray(transl[k]), # (1, 3) + "scores": CoreMLArray(np.array([sc], dtype=np.float32)), + "shape": CoreMLArray(betas[k]), # (10,) + "expression": CoreMLArray(expr[k]), # (10,) + }) + return humans diff --git a/data_only_viz/scripts/capture_actions.py b/data_only_viz/scripts/capture_actions.py new file mode 100644 index 0000000..331858d --- /dev/null +++ b/data_only_viz/scripts/capture_actions.py @@ -0,0 +1,74 @@ +"""Record webcam frames + timestamps for action-head training. + +Usage: + uv run python -m data_only_viz.scripts.capture_actions \ + --session sess03 --duration 600 +""" +from __future__ import annotations + +import argparse +import logging +import time +from pathlib import Path + +import cv2 + +LOG = logging.getLogger("capture_actions") +RAW_DIR = Path("~/.cache/av-live-action/raw").expanduser() + + +def capture(session: str, duration_s: float, + cam_index: int = 0, fps: int = 30, + size: int = 672) -> Path: + RAW_DIR.mkdir(parents=True, exist_ok=True) + out = RAW_DIR / f"{session}.mp4" + ts_out = RAW_DIR / f"{session}.ts.txt" + cap = cv2.VideoCapture(cam_index) + if not cap.isOpened(): + raise RuntimeError(f"cannot open camera {cam_index}") + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter(str(out), fourcc, fps, (size, size)) + try: + t_start = time.perf_counter() + with ts_out.open("w") as ts_f: + n = 0 + while time.perf_counter() - t_start < duration_s: + ok, frame = cap.read() + if not ok: + LOG.warning("frame read failed") + break + h, w = frame.shape[:2] + side = min(h, w) + y0 = (h - side) // 2 + x0 = (w - side) // 2 + crop = frame[y0:y0 + side, x0:x0 + side] + resized = cv2.resize(crop, (size, size)) + writer.write(resized) + ts_f.write(f"{n} {time.perf_counter() - t_start:.6f}\n") + n += 1 + cv2.imshow("capture (q=quit)", resized) + if cv2.waitKey(1) & 0xFF == ord("q"): + break + LOG.info("wrote %s (%d frames)", out, n) + return out + finally: + cap.release() + writer.release() + cv2.destroyAllWindows() + + +def _cli() -> None: + p = argparse.ArgumentParser() + p.add_argument("--session", required=True) + p.add_argument("--duration", type=float, default=600.0) + p.add_argument("--cam-index", type=int, default=0) + p.add_argument("--fps", type=int, default=30) + args = p.parse_args() + logging.basicConfig(level=logging.INFO, + format="%(asctime)s [%(name)s] %(message)s") + capture(args.session, args.duration, + cam_index=args.cam_index, fps=args.fps) + + +if __name__ == "__main__": + _cli() diff --git a/data_only_viz/scripts/coreml_full_probe.py b/data_only_viz/scripts/coreml_full_probe.py index 241ce15..f1f4387 100644 --- a/data_only_viz/scripts/coreml_full_probe.py +++ b/data_only_viz/scripts/coreml_full_probe.py @@ -157,6 +157,32 @@ if hasattr(model.backbone, "encoder") and hasattr(model.backbone.encoder, # torch.inverse(K) plante coremltools (op non implementee). Comme K est # fixe (camera intrinsics avec focal=IMG_SIZE), on pre-calcule K_inv # en closed-form et on l'utilise comme buffer module-level. +print("==> Patching roma.rotmat_to_rotvec (branchless atan2)") +# roma.rotmat_to_rotvec utilise torch.empty + 8 index_put_ qui se +# traduisent en CoreML par scatter_nd successifs sur un buffer +# garbage-initialise. Resultat : cellules non touchees restent NaN, +# propagees via quat normalization -> v3d/transl all-NaN. +# Remplacement branchless via atan2 : pas de torch.empty, pas +# d'index_put_, juste des stack/clamp/norm/atan2 stables CoreML. +# Precision vs roma original : 2.26e-6 L_inf sur batch random. +import roma as _roma + +def _rotmat_to_rotvec_branchless(R, eps=1e-6): + w = torch.stack([ + R[..., 2, 1] - R[..., 1, 2], + R[..., 0, 2] - R[..., 2, 0], + R[..., 1, 0] - R[..., 0, 1], + ], dim=-1) * 0.5 + trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] + cos_theta = ((trace - 1.0) * 0.5).clamp(-1.0, 1.0) + sin_theta = torch.norm(w, dim=-1) + theta = torch.atan2(sin_theta, cos_theta) + sin_theta_safe = sin_theta.clamp(min=eps) + return w * (theta / sin_theta_safe).unsqueeze(-1) + +_roma.rotmat_to_rotvec = _rotmat_to_rotvec_branchless + + print("==> Patching utils.camera.inverse_perspective_projection") import utils.camera as _camera @@ -170,11 +196,28 @@ _K_INV_PRE = torch.tensor([ ]) def inverse_perspective_projection_fixed(points, K, distance): - """Bypass torch.inverse : utilise K_inv pre-calcule en closed-form - (notre K est connu et fixe). Le K argument est ignore.""" - K_inv = _K_INV_PRE.to(points.device).to(points.dtype) - points = torch.cat([points, torch.ones_like(points[..., :1])], -1) - points = torch.einsum('bij,bkj->bki', K_inv, points) + """Bypass torch.inverse + einsum + matmul pour eviter le bug + coremltools de broadcast batch 1->K sur ces ops. K_inv etant + fixe et structure (diag + translate), on ecrit les composantes + explicitement en ops elementaires. + + K_inv = [[1/f, 0, -cx/f], [0, 1/f, -cy/f], [0, 0, 1]] + Pour points (b, N, 3) : out = points @ K_inv.T donne : + out[..., 0] = points[..., 0]/f - (cx/f) * points[..., 2] + out[..., 1] = points[..., 1]/f - (cy/f) * points[..., 2] + out[..., 2] = points[..., 2] + """ + points_hom = torch.cat([points, torch.ones_like(points[..., :1])], -1) + inv_f = 1.0 / focal_val + cx_over_f = cx / focal_val + cy_over_f = cy / focal_val + x = points_hom[..., 0:1] + y = points_hom[..., 1:2] + z = points_hom[..., 2:3] + out0 = x * inv_f - z * cx_over_f + out1 = y * inv_f - z * cy_over_f + out2 = z + points = torch.cat([out0, out1, out2], dim=-1) if distance is None: return points points = points * distance @@ -190,6 +233,26 @@ model_mod.inverse_perspective_projection = inverse_perspective_projection_fixed import blocks.smpl_layer as _smpl_layer _smpl_layer.inverse_perspective_projection = inverse_perspective_projection_fixed +# Aussi perspective_projection (utilise dans smpl_layer.py:143-144 pour +# j2d et v2d) -> rewrite einsum en matmul pour le meme broadcast bug. +def perspective_projection_fixed(x, K): + """Element-wise rewrite de la projection perspective avec K fixe + (focal=IMG_SIZE, cx=cy=IMG_SIZE/2). Bypass matmul/einsum pour eviter + les bugs broadcast coremltools. + K = [[f, 0, cx], [0, f, cy], [0, 0, 1]] + out[..., 0] = f * x_norm + cx * z_norm (mais on veut [..., :2]) + = f * (x/z) + cx + out[..., 1] = f * (y/z) + cy + """ + z = x[..., 2:3] + px = x[..., 0:1] / z * focal_val + cx + py = x[..., 1:2] / z * focal_val + cy + return torch.cat([px, py], dim=-1) + +_camera.perspective_projection = perspective_projection_fixed +_utils_pkg.perspective_projection = perspective_projection_fixed +_smpl_layer.perspective_projection = perspective_projection_fixed + # === Wrapper qui produit tuple fixe === class TracedMHMR(nn.Module): @@ -222,6 +285,12 @@ class TracedMHMR(nn.Module): ]).squeeze(-1) shape = torch.stack([h["shape"] for h in humans]) expr = torch.stack([h["expression"] for h in humans]) + # NOTE: CoreML mlprogram conversion currently produces all-NaN + # outputs for v3d and transl while PyTorch eager produces valid + # finite values from the same trace. nan_to_num here masks the + # symptom but yields all-zero meshes (no information). Leave + # raw outputs and let downstream decide; investigation tracked + # in task #2 (op-by-op bisection needed). return v3d, transl, scores, shape, expr @@ -422,6 +491,28 @@ def _diagonal_general(context, node): _TORCH_OPS_REGISTRY.name_to_func_mapping["diagonal"] = _diagonal_general + +# Instrument reshape pour logger node source au moment de l'erreur. +from coremltools.converters.mil.mil.ops.defs.iOS15 import tensor_transformation as _tt +_orig_reshape_ti = _tt.reshape.type_inference + + +def _reshape_ti_logged(self): + try: + return _orig_reshape_ti(self) + except ValueError as e: + if "Invalid target shape" in str(e): + try: + from_shape = list(self.x.shape) + target = list(self.shape.val) if hasattr(self.shape, "val") else "?" + print(f" >>> RESHAPE FAIL : name={self.name} from={from_shape} target={target}") + except Exception: + pass + raise + + +_tt.reshape.type_inference = _reshape_ti_logged + try: mlmodel = ct.convert( traced, @@ -433,6 +524,10 @@ try: compute_units=ct.ComputeUnit.CPU_AND_GPU, minimum_deployment_target=ct.target.macOS15, convert_to="mlprogram", + # FP16 OK depuis le patch roma branchless (cf rapport bisection + # 2026-05-13) : la source du NaN etait torch.empty + index_put_ + # dans roma.rotmat_to_rotvec, pas la precision. + compute_precision=ct.precision.FLOAT16, ) out_path = "/tmp/multihmr_full_672_s.mlpackage" mlmodel.save(out_path) diff --git a/data_only_viz/scripts/extract_j3d_offline.py b/data_only_viz/scripts/extract_j3d_offline.py index 1ad9254..6fb5b09 100644 --- a/data_only_viz/scripts/extract_j3d_offline.py +++ b/data_only_viz/scripts/extract_j3d_offline.py @@ -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, HANDS_KP_DIMS, HANDS_KP_TOTAL +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,14 +52,37 @@ 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 + # CoreMLArray wraps numpy but lacks __array__; unwrap before asarray. + if hasattr(v3d, "numpy") and not isinstance(v3d, np.ndarray): + v3d = v3d.numpy() 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, @@ -86,14 +114,20 @@ 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, + "hands_kp": np.zeros( + (HANDS_KP_TOTAL, HANDS_KP_DIMS), dtype=np.float32 + ).tolist(), }) + "\n") n_rows += 1 n_frames += 1 diff --git a/data_only_viz/scripts/extract_mediapipe_offline.py b/data_only_viz/scripts/extract_mediapipe_offline.py new file mode 100644 index 0000000..18fc108 --- /dev/null +++ b/data_only_viz/scripts/extract_mediapipe_offline.py @@ -0,0 +1,208 @@ +"""Extract action-head v3 jsonl rows from a recorded MP4 using MediaPipe +Holistic. Populates real hands_kp (42, 3) and mouth_open (face lips +distance), unlike extract_j3d_offline.py (SMPL-X path) which writes zeros +for hands_kp. + +Output jsonl row format (matches dataset.py load_frames_jsonl) : + + { + "ts": float seconds, + "session": str, + "pid": int (always 0 — Holistic is single-person), + "j3d": [[32, 3]] floats (body22 + 10 fingertips), + "expression": [10] zeros (MediaPipe has no SMPL-X PCA), + "mouth_open": float (lips inner distance), + "hands_kp": [[42, 3]] floats (21 L + 21 R, zero-padded if absent), + } + +Usage : + + uv run python -m data_only_viz.scripts.extract_mediapipe_offline \ + --session sess03 \ + --video ~/.cache/av-live-action/raw/sess03.mp4 \ + --out ~/.cache/av-live-action/raw/sess03_mp.jsonl +""" +from __future__ import annotations + +import argparse +import json +import logging +from pathlib import Path + +import cv2 +import numpy as np + +from data_only_viz.action_head import ( + EXPR_DIM, + HANDS_KP_DIMS, + HANDS_KP_PER_HAND, + HANDS_KP_TOTAL, + J3D_BODY, + J3D_FINGERS, + J3D_FINGERS_PER_HAND, + J3D_JOINTS, +) +from data_only_viz.action_head_pub import ( + MEDIAPIPE_HAND_FINGERTIPS, + MEDIAPIPE_LIP_LOWER_INNER, + MEDIAPIPE_LIP_UPPER_INNER, + MEDIAPIPE_TO_22, +) + +LOG = logging.getLogger("extract_mediapipe_offline") +DEFAULT_OUT_DIR = Path("~/.cache/av-live-action/raw").expanduser() + + +def _build_landmarker(): + """Build a MediaPipe HolisticLandmarker in VIDEO running mode.""" + from mediapipe.tasks.python import vision + from mediapipe.tasks.python.core.base_options import BaseOptions + from data_only_viz.holistic import _ensure_model + model_path = _ensure_model() + opts = vision.HolisticLandmarkerOptions( + base_options=BaseOptions(model_asset_path=str(model_path)), + running_mode=vision.RunningMode.VIDEO, + min_pose_detection_confidence=0.3, + min_pose_landmarks_confidence=0.3, + min_face_detection_confidence=0.3, + min_face_landmarks_confidence=0.3, + min_hand_landmarks_confidence=0.3, + ) + return vision.HolisticLandmarker.create_from_options(opts) + + +def _lmk_list_to_array(lmks) -> np.ndarray | None: + """Convert MediaPipe NormalizedLandmark / Landmark list to (N, 3) array.""" + if lmks is None: + return None + try: + return np.asarray( + [(lm.x, lm.y, getattr(lm, "z", 0.0)) for lm in lmks], + dtype=np.float32, + ) + except (AttributeError, TypeError): + return None + + +def _build_j3d32(body3d_arr: np.ndarray | None, + hands_kp42: np.ndarray) -> np.ndarray | None: + """Map MediaPipe body3d (33, 3) + hands_kp (42, 3) -> j3d (32, 3). + + body22 indices via MEDIAPIPE_TO_22, fingertips from hands_kp idx + MEDIAPIPE_HAND_FINGERTIPS (4, 8, 12, 16, 20) for each side. + """ + if body3d_arr is None or body3d_arr.shape[0] < 33: + return None + body22 = body3d_arr[list(MEDIAPIPE_TO_22)].astype(np.float32) + tips = np.zeros((J3D_FINGERS, 3), dtype=np.float32) + for side_idx in (0, 1): + base = side_idx * HANDS_KP_PER_HAND + for k, mp_tip in enumerate(MEDIAPIPE_HAND_FINGERTIPS): + if base + mp_tip < hands_kp42.shape[0]: + tips[side_idx * J3D_FINGERS_PER_HAND + k] = hands_kp42[base + mp_tip] + return np.concatenate([body22, tips], axis=0) + + +def _mouth_open(face_arr: np.ndarray | None) -> float: + if face_arr is None or face_arr.shape[0] <= MEDIAPIPE_LIP_LOWER_INNER: + return 0.0 + upper = face_arr[MEDIAPIPE_LIP_UPPER_INNER] + lower = face_arr[MEDIAPIPE_LIP_LOWER_INNER] + return float(np.linalg.norm(upper - lower)) + + +def _hands_kp42(left_arr: np.ndarray | None, + right_arr: np.ndarray | None) -> np.ndarray: + out = np.zeros((HANDS_KP_TOTAL, HANDS_KP_DIMS), dtype=np.float32) + if left_arr is not None and left_arr.shape[0] >= HANDS_KP_PER_HAND: + out[:HANDS_KP_PER_HAND] = left_arr[:HANDS_KP_PER_HAND] + if right_arr is not None and right_arr.shape[0] >= HANDS_KP_PER_HAND: + out[HANDS_KP_PER_HAND:] = right_arr[:HANDS_KP_PER_HAND] + return out + + +def extract(session: str, video: Path, out: Path) -> int: + """Run MediaPipe Holistic on every frame of video, write jsonl rows. + + Returns the number of frames where at least body3d was detected + (rows written). Frames with no person are silently skipped. + """ + import mediapipe as mp + out.parent.mkdir(parents=True, exist_ok=True) + cap = cv2.VideoCapture(str(video)) + if not cap.isOpened(): + raise RuntimeError(f"cannot open {video}") + fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 + landmarker = _build_landmarker() + n_frames = 0 + n_rows = 0 + expr_zeros_list = np.zeros(EXPR_DIM, dtype=np.float32).tolist() + try: + with out.open("w") as f: + while True: + ok, frame = cap.read() + if not ok: + break + rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb) + ts_ms = int(n_frames * 1000 / fps) + try: + res = landmarker.detect_for_video(mp_img, ts_ms) + except Exception: + LOG.exception("detect failed at frame=%d", n_frames) + n_frames += 1 + continue + body3d = _lmk_list_to_array( + getattr(res, "pose_world_landmarks", None) + ) + face_arr = _lmk_list_to_array( + getattr(res, "face_landmarks", None) + ) + left_arr = _lmk_list_to_array( + getattr(res, "left_hand_landmarks", None) + ) + right_arr = _lmk_list_to_array( + getattr(res, "right_hand_landmarks", None) + ) + hands_kp42 = _hands_kp42(left_arr, right_arr) + j3d32 = _build_j3d32(body3d, hands_kp42) + if j3d32 is None: + n_frames += 1 + continue + ts = n_frames / fps + mouth = _mouth_open(face_arr) + f.write(json.dumps({ + "ts": ts, + "session": session, + "pid": 0, + "j3d": j3d32.tolist(), + "expression": expr_zeros_list, + "mouth_open": mouth, + "hands_kp": hands_kp42.tolist(), + }) + "\n") + n_rows += 1 + n_frames += 1 + if n_frames % 100 == 0: + LOG.info("frame=%d rows=%d", n_frames, n_rows) + finally: + cap.release() + landmarker.close() + LOG.info("done : %d frames, %d rows -> %s", n_frames, n_rows, out) + return n_rows + + +def _cli() -> None: + p = argparse.ArgumentParser() + p.add_argument("--session", required=True) + p.add_argument("--video", required=True, type=Path) + p.add_argument("--out", type=Path) + args = p.parse_args() + logging.basicConfig(level=logging.INFO, + format="%(asctime)s [%(name)s] %(message)s") + DEFAULT_OUT_DIR.mkdir(parents=True, exist_ok=True) + out = args.out or (DEFAULT_OUT_DIR / f"{args.session}_mp.jsonl") + extract(args.session, args.video, out) + + +if __name__ == "__main__": + _cli() diff --git a/data_only_viz/scripts/setup_multihmr.sh b/data_only_viz/scripts/setup_multihmr.sh index b51df54..fcf96cb 100755 --- a/data_only_viz/scripts/setup_multihmr.sh +++ b/data_only_viz/scripts/setup_multihmr.sh @@ -49,4 +49,71 @@ if [ ! -e "$CACHE/multi-hmr/models" ]; then ln -sfn ../models "$CACHE/multi-hmr/models" fi +# CoreML conversion patches : remplace les torch.einsum dans utils/camera.py +# par des ops element-wise (broadcast-friendly). Sans ca, ct.convert echoue +# avec "Invalid target shape in reshape op ([1, N, 3] to [K*N, 3, 1])" +# quand batch K detections != 1. Idempotent. +CAM="$CACHE/multi-hmr/utils/camera.py" +if [ -f "$CAM" ] && ! grep -q "_apply_intrinsics_componentwise" "$CAM"; then + echo "==> Patch utils/camera.py (einsum -> componentwise)" + python3 - "$CAM" <<'PYEOF' +import sys, pathlib +p = pathlib.Path(sys.argv[1]) +src = p.read_text() +helper = ''' +def _apply_intrinsics_componentwise(K, y): + """CoreML-friendly: out[b,k,i] = sum_j K[b,i,j] * y[b,k,j] + Replaces torch.einsum('bij,bkj->bki', K, y) with pure broadcast ops. + """ + K00 = K[:, 0:1, 0:1]; K01 = K[:, 0:1, 1:2]; K02 = K[:, 0:1, 2:3] + K10 = K[:, 1:2, 0:1]; K11 = K[:, 1:2, 1:2]; K12 = K[:, 1:2, 2:3] + K20 = K[:, 2:3, 0:1]; K21 = K[:, 2:3, 1:2]; K22 = K[:, 2:3, 2:3] + y0 = y[:, :, 0:1]; y1 = y[:, :, 1:2]; y2 = y[:, :, 2:3] + out0 = K00 * y0 + K01 * y1 + K02 * y2 + out1 = K10 * y0 + K11 * y1 + K12 * y2 + out2 = K20 * y0 + K21 * y1 + K22 * y2 + return torch.cat([out0, out1, out2], dim=-1) + + +''' +src = src.replace( + "def perspective_projection(x, K):", + helper + "def perspective_projection(x, K):", +) +src = src.replace( + "y = torch.einsum('bij,bkj->bki', K, y) # (bs, N, 3)", + "y = _apply_intrinsics_componentwise(K, y)", +) +src = src.replace( + "points = torch.einsum('bij,bkj->bki', torch.inverse(K), points)", + "points = _apply_intrinsics_componentwise(torch.inverse(K), points)", +) +p.write_text(src) +print(" camera.py patched") +PYEOF +fi + +# CoreML conversion patch : smplx/lbs.py landmarks einsum (mеme bug broadcast) +# Patch best-effort sur tous les venvs presents (data_only_viz + /tmp/coreml312). +for VENV in \ + "$(dirname "$(dirname "$(readlink -f "$0")")")/.venv" \ + "/tmp/coreml312"; do + LBS="$VENV/lib/python3.14/site-packages/smplx/lbs.py" + [ -f "$LBS" ] || LBS="$VENV/lib/python3.12/site-packages/smplx/lbs.py" + if [ -f "$LBS" ] && grep -q "torch.einsum('blfi,blf->bli'" "$LBS"; then + echo "==> Patch $LBS (landmarks einsum)" + python3 - "$LBS" <<'PYEOF' +import sys, pathlib +p = pathlib.Path(sys.argv[1]) +s = p.read_text() +s = s.replace( + "landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])\n return landmarks", + "# CoreML-friendly: replace einsum('blfi,blf->bli', ...) with broadcast+sum\n landmarks = (lmk_vertices * lmk_bary_coords.unsqueeze(-1)).sum(dim=2)\n return landmarks", +) +p.write_text(s) +print(" smplx/lbs.py patched") +PYEOF + fi +done + echo "Setup OK. Cache : $CACHE" diff --git a/data_only_viz/scripts/train_on_studio.sh b/data_only_viz/scripts/train_on_studio.sh new file mode 100755 index 0000000..e4ec3fe --- /dev/null +++ b/data_only_viz/scripts/train_on_studio.sh @@ -0,0 +1,111 @@ +#!/usr/bin/env bash +# Train action-head on MacStudio M3 Ultra (Tailscale 100.116.92.12). +# +# SSH direct grosmac→studio is broken since reboot 2026-05-12 ; +# we route via electron-server bastion (cf. CLAUDE.md root). +# +# Usage: +# ./train_on_studio.sh # uses defaults +# ./train_on_studio.sh --epochs 80 --lr 5e-4 +# +# Local layout : +# ~/.cache/av-live-action/dataset/dataset.jsonl (input) +# ~/.cache/av-live-action/checkpoints/ (output, after rsync back) +# +# Remote layout : +# studio:~/av-live-action/repo/ (rsynced code subset) +# studio:~/av-live-action/dataset/ (rsynced dataset) +# studio:~/av-live-action/checkpoints/ (training output) + +set -euo pipefail + +BASTION_USER_HOST="${BASTION_USER_HOST:-electron-server}" +STUDIO_USER_HOST="${STUDIO_USER_HOST:-clems@100.116.92.12}" +STUDIO_USER="${STUDIO_USER:-clems}" +STUDIO_UV="${STUDIO_UV:-/opt/homebrew/bin/uv}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")"/../.. && pwd)" +LOCAL_CACHE="$HOME/.cache/av-live-action" +LOCAL_DATASET="$LOCAL_CACHE/dataset" +LOCAL_CKPT="$LOCAL_CACHE/checkpoints" + +REMOTE_ROOT="/Users/${STUDIO_USER}/av-live-action" +REMOTE_REPO="$REMOTE_ROOT/repo" +REMOTE_DATASET="$REMOTE_ROOT/dataset" +REMOTE_CKPT="$REMOTE_ROOT/checkpoints" + +DATASET_FILE="${DATASET_FILE:-$LOCAL_DATASET/dataset.jsonl}" +CKPT_NAME="${CKPT_NAME:-action_head.pt}" + +# Quote train args defensively before forwarding through bastion ssh + +# studio ssh (each layer reparses). Reject single quotes — they break +# the single-quoted payload in bastion_ssh and could allow injection. +for a in "$@"; do + if [[ "$a" == *"'"* ]]; then + printf '[train_on_studio] forbidden single quote in arg: %s\n' "$a" >&2 + exit 3 + fi +done +TRAIN_ARGS="$(printf '%q ' "$@")" + +log() { printf '[train_on_studio] %s\n' "$*" >&2; } + +[[ -f "$DATASET_FILE" ]] || { log "missing dataset: $DATASET_FILE"; exit 2; } +mkdir -p "$LOCAL_CKPT" + +bastion_ssh() { + # The remote shell on the bastion must receive the studio command + # as a single argument, otherwise `;` and `&&` are parsed + # bastion-side instead of studio-side. + # All paths in commands MUST be absolute (no $HOME, no ~) since + # we use single-quotes for the studio-side payload. + ssh -o ConnectTimeout=5 "$BASTION_USER_HOST" \ + "ssh -o ConnectTimeout=5 $STUDIO_USER_HOST '$*'" +} + +bastion_rsync() { + # rsync via ssh ProxyJump through bastion. Direct grosmac->studio + # known_hosts entry may be stale (SSH direct broken since reboot + # 2026-05-12). accept-new lets us add the key on first use. + local src="$1" dst="$2" + rsync -avz --delete \ + -e "ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new -A -J $BASTION_USER_HOST" \ + "$src" "$dst" +} + +log "== Studio reachability ==" +bastion_ssh "echo studio OK ; $STUDIO_UV --version" + +log "== Push code subset ==" +bastion_ssh "mkdir -p $REMOTE_REPO/data_only_viz $REMOTE_DATASET $REMOTE_CKPT" +rsync -avz --delete \ + --exclude='.venv/' --exclude='__pycache__/' --exclude='.pytest_cache/' \ + --exclude='.ruff_cache/' --exclude='*.pyc' --exclude='.DS_Store' \ + --exclude='web/' --exclude='shaders/' \ + -e "ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new -A -J $BASTION_USER_HOST" \ + "$REPO_ROOT/data_only_viz/" \ + "$STUDIO_USER_HOST:av-live-action/repo/data_only_viz/" + +log "== Push dataset ==" +bastion_rsync "$LOCAL_DATASET/" "$STUDIO_USER_HOST:av-live-action/dataset/" + +log "== Remote uv sync ==" +# multihmr extra pulls torch (action-head training needs torch but no pyobjc). +# We piggy-back on the multihmr extras since torch is the main thing we need. +bastion_ssh "cd $REMOTE_REPO && $STUDIO_UV sync --no-progress --project data_only_viz --extra multihmr" + +log "== Remote train (MPS) ==" +# cwd must be the PARENT of data_only_viz/ so the package is importable as +# top-level. uv resolves the env via --project data_only_viz. +bastion_ssh "cd $REMOTE_REPO && \ + $STUDIO_UV run --project data_only_viz python -m data_only_viz.training.train_action_head \ + --dataset $REMOTE_DATASET/$(basename "$DATASET_FILE") \ + --ckpt-out $REMOTE_CKPT/$CKPT_NAME \ + --device mps \ + $TRAIN_ARGS" + +log "== Pull checkpoint back ==" +bastion_rsync "$STUDIO_USER_HOST:av-live-action/checkpoints/" "$LOCAL_CKPT/" + +log "== Done. Checkpoint: $LOCAL_CKPT/$CKPT_NAME ==" +ls -la "$LOCAL_CKPT/$CKPT_NAME" diff --git a/data_only_viz/smplx_osc_sender.py b/data_only_viz/smplx_osc_sender.py index aa7b379..18dbdd8 100644 --- a/data_only_viz/smplx_osc_sender.py +++ b/data_only_viz/smplx_osc_sender.py @@ -25,6 +25,7 @@ from typing import Sequence import numpy as np +from .mesh_rigger import MeshRigger from .state import SMPLXPerson, State LOG = logging.getLogger("smplx_tcp") @@ -35,7 +36,8 @@ PORT = 57130 class SMPLXTCPSender: def __init__(self, state: State, host: str = "127.0.0.1", - port: int = PORT, target_fps: float = 12.0) -> None: + port: int = PORT, target_fps: float = 30.0, + enable_rigging: bool = True) -> None: self.state = state self.host = host self.port = port @@ -43,6 +45,9 @@ class SMPLXTCPSender: self._stop = threading.Event() self._thread: threading.Thread | None = None self._sock: socket.socket | None = None + # Hybrid keyframe rigging : entre deux keyframes Multi-HMR (~3 fps), + # on translate le mesh via le delta pelvis Apple Vision (30 fps). + self._rigger = MeshRigger(state) if enable_rigging else None def start(self) -> None: self._thread = threading.Thread( @@ -124,6 +129,9 @@ class SMPLXTCPSender: def _run(self) -> None: last_warn = 0.0 + n_sent = 0 + n_rigged = 0 + next_hb = time.monotonic() + 5.0 while not self._stop.is_set(): t0 = time.monotonic() if not self._ensure_connected(): @@ -136,8 +144,30 @@ class SMPLXTCPSender: with self.state.lock(): persons = list(self.state.persons_smplx) + body_kp = list(self.state.persons_body) if hasattr( + self.state, "persons_body") else [] + body_ids = list(self.state.persons_body_ids) if hasattr( + self.state, "persons_body_ids") else ( + list(range(len(body_kp))) if body_kp else []) + + if persons and self._rigger is not None: + rigged = self._rigger.apply( + persons, body_kp, body_ids, t0) + if rigged is not persons: + n_rigged += 1 + persons = rigged + + if t0 >= next_hb: + fps = n_sent / 5.0 + rig_pct = (n_rigged / n_sent * 100.0) if n_sent else 0.0 + LOG.info("hb: %.1f fps tcp, %.0f%% rigged", + fps, rig_pct) + n_sent = 0 + n_rigged = 0 + next_hb = t0 + 5.0 if persons: + n_sent += 1 t_ser_start = time.monotonic() payload = self._serialize_persons(persons) t_send_start = time.monotonic() diff --git a/data_only_viz/tests/test_action_head_features.py b/data_only_viz/tests/test_action_head_features.py new file mode 100644 index 0000000..b6df78d --- /dev/null +++ b/data_only_viz/tests/test_action_head_features.py @@ -0,0 +1,116 @@ +"""Unit tests for ActionHead feature extraction and buffers.""" +from __future__ import annotations + +import numpy as np +import pytest + + +def test_module_imports() -> None: + from data_only_viz import action_head + assert hasattr(action_head, "FeatureExtractor") + assert hasattr(action_head, "PerPersonBuffer") + assert hasattr(action_head, "ActionHead") + assert action_head.WINDOW_LEN == 16 + assert action_head.J3D_JOINTS == 32 + assert action_head.FEATURE_DIM == 428 + assert action_head.HANDS_KP_TOTAL == 42 + assert action_head.HANDS_KP_FLAT == 126 + 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=(32, 3)).astype(np.float32) + + +def test_buffer_starts_empty() -> None: + from data_only_viz.action_head import PerPersonBuffer + buf = PerPersonBuffer() + assert len(buf) == 0 + assert buf.frames_for(7) == [] + + +def test_buffer_append_grows_per_pid() -> None: + from data_only_viz.action_head import PerPersonBuffer + buf = PerPersonBuffer() + buf.append(pid=1, j3d=_rand_j3d(1)) + buf.append(pid=1, j3d=_rand_j3d(2)) + buf.append(pid=2, j3d=_rand_j3d(3)) + assert len(buf.frames_for(1)) == 2 + assert len(buf.frames_for(2)) == 1 + + +def test_buffer_max_len_16() -> None: + from data_only_viz.action_head import PerPersonBuffer, WINDOW_LEN + buf = PerPersonBuffer() + for i in range(WINDOW_LEN + 5): + buf.append(pid=1, j3d=_rand_j3d(i)) + assert len(buf.frames_for(1)) == WINDOW_LEN + + +def test_buffer_forget_releases_pid() -> None: + from data_only_viz.action_head import PerPersonBuffer + buf = PerPersonBuffer() + buf.append(pid=1, j3d=_rand_j3d(0)) + buf.forget(1) + assert buf.frames_for(1) == [] + assert len(buf) == 0 + + +def test_buffer_rejects_bad_shape() -> None: + from data_only_viz.action_head import PerPersonBuffer + buf = PerPersonBuffer() + 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 == (428,) + assert feat.dtype == np.float32 + assert not np.isnan(feat).any() + + +def test_feature_extractor_short_buffer_pads() -> None: + from data_only_viz.action_head import FeatureExtractor, FEATURE_DIM + frames = [_rand_j3d(0), _rand_j3d(1), _rand_j3d(2)] + feat = FeatureExtractor.from_buffer(frames) + assert feat.shape == (FEATURE_DIM,) + + +def test_feature_extractor_static_buffer_zero_velocity() -> None: + from data_only_viz.action_head import FeatureExtractor, WINDOW_LEN, J3D_JOINTS + static = _rand_j3d(42) + frames = [static.copy() for _ in range(WINDOW_LEN)] + feat = FeatureExtractor.from_buffer(frames) + vel_block = feat[J3D_JOINTS * 3 : J3D_JOINTS * 3 * 2] + assert np.allclose(vel_block, 0.0, atol=1e-6) + + +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((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 / 32) < 1e-4 + assert abs(kin[1]) < 1e-4 + + +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((32, 3), dtype=np.float32) + f[WRIST_LEFT, 0] = 0.05 * t + f[WRIST_RIGHT, 0] = -0.05 * t + frames.append(f) + kin = FeatureExtractor.kinetics(frames) + assert kin[2] > 0.9 diff --git a/data_only_viz/tests/test_action_head_model.py b/data_only_viz/tests/test_action_head_model.py new file mode 100644 index 0000000..939d8c4 --- /dev/null +++ b/data_only_viz/tests/test_action_head_model.py @@ -0,0 +1,83 @@ +"""Tests for ActionHead model (forward, step, checkpoint roundtrip).""" +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + + +def _rand_j3d(seed: int = 0) -> np.ndarray: + rng = np.random.default_rng(seed) + return rng.normal(size=(32, 3)).astype(np.float32) + + +def test_model_forward_shape() -> None: + from data_only_viz.action_head import ActionHeadModel, FEATURE_DIM, NUM_CLASSES + model = ActionHeadModel() + x = torch.zeros(1, FEATURE_DIM) + h = model.init_hidden(batch=1) + logits, h_new = model(x, h) + assert logits.shape == (1, NUM_CLASSES) + assert h_new.shape == h.shape + + +def test_model_param_count_under_100k() -> None: + from data_only_viz.action_head import ActionHeadModel + model = ActionHeadModel() + n = sum(p.numel() for p in model.parameters()) + assert n < 100_000, f"too many params: {n}" + + +def test_action_head_step_warmup_returns_debout() -> None: + from data_only_viz.action_head import ActionHead, LABELS + head = ActionHead(ckpt_path=None) + label, probs, kin = head.step(pid=1, j3d=_rand_j3d(0)) + assert label == LABELS[0] + assert probs.shape == (3,) + assert pytest.approx(float(probs[0]), abs=1e-6) == 1.0 + assert kin.shape == (3,) + assert float(kin[0]) == 0.0 + + +def test_action_head_step_after_warmup_returns_some_label() -> None: + from data_only_viz.action_head import ActionHead, LABELS + head = ActionHead(ckpt_path=None) + for i in range(5): + label, probs, kin = head.step(pid=1, j3d=_rand_j3d(i)) + assert label in LABELS + assert abs(float(probs.sum()) - 1.0) < 1e-5 + + +def test_action_head_forget_resets_hidden_state(tmp_path: Path) -> None: + from data_only_viz.action_head import ActionHead + head = ActionHead(ckpt_path=None) + for i in range(5): + head.step(pid=1, j3d=_rand_j3d(i)) + assert 1 in head._hidden + head.forget(1) + assert 1 not in head._hidden + assert head._buffers.frames_for(1) == [] + + +def test_action_head_checkpoint_roundtrip(tmp_path: Path) -> None: + from data_only_viz.action_head import ActionHead, ActionHeadModel + model = ActionHeadModel() + ckpt = tmp_path / "ah.pt" + torch.save({"model_state_dict": model.state_dict(), + "version": 1}, ckpt) + head = ActionHead(ckpt_path=ckpt) + for k, v in head._model.state_dict().items(): + assert torch.allclose(v, model.state_dict()[k]) + + +def test_action_head_step_handles_nan() -> None: + from data_only_viz.action_head import ActionHead, LABELS + head = ActionHead(ckpt_path=None) + j = _rand_j3d(0) + j[5, 1] = float("nan") + label, probs, _kin = head.step(pid=1, j3d=j) + assert label in LABELS + assert not np.isnan(probs).any() diff --git a/data_only_viz/tests/test_action_head_pub.py b/data_only_viz/tests/test_action_head_pub.py new file mode 100644 index 0000000..540e5ad --- /dev/null +++ b/data_only_viz/tests/test_action_head_pub.py @@ -0,0 +1,154 @@ +"""Tests for ActionHeadPublisher.""" +from __future__ import annotations + +import threading +from unittest.mock import MagicMock + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + + +class _FakeState: + def __init__(self) -> None: + self.persons_smplx = [] + self.smplx_last_t = 0.0 + self.persons_body3d = [] + self.persons_body_ids = [] + self.pose_last_t = 0.0 + self.persons_hands = [] + self.persons_hands_ids = [] + self.persons_face = [] + self.persons_face_ids = [] + self._lock = threading.RLock() + + def lock(self): + return self._lock + + +def _make_smplx_person(pid: int, seed: int = 0) -> dict: + rng = np.random.default_rng(seed) + return {"pid": pid, "v3d": rng.normal(size=(10475, 3)).astype(np.float32)} + + +def test_publisher_smplx_source_emits_osc() -> None: + from data_only_viz.action_head_pub import ActionHeadPublisher + state = _FakeState() + bridge = MagicMock() + pub = ActionHeadPublisher(state, bridge, ckpt_path=None) + state.persons_smplx = [_make_smplx_person(7)] + state.smplx_last_t = 1.0 + pub._tick(t_now=0.0) + actions = [c for c in bridge.send_action.call_args_list] + assert len(actions) == 1 + assert actions[0].kwargs.get("pid", actions[0].args[0]) == 7 + bridge.send_enter.assert_called_with(pid=7) + + +def test_publisher_falls_back_to_mediapipe_body3d() -> None: + from data_only_viz.action_head_pub import ActionHeadPublisher + state = _FakeState() + bridge = MagicMock() + pub = ActionHeadPublisher(state, bridge, ckpt_path=None) + state.persons_body3d = [[(0.1 * i, 0.2 * i, 0.3 * i) for i in range(33)]] + state.persons_body_ids = [42] + state.pose_last_t = 1.0 + pub._tick(t_now=0.0) + bridge.send_action.assert_called_once() + bridge.send_enter.assert_called_with(pid=42) + + +def test_publisher_purges_lost_pid() -> None: + from data_only_viz.action_head_pub import ActionHeadPublisher + state = _FakeState() + bridge = MagicMock() + pub = ActionHeadPublisher(state, bridge, ckpt_path=None) + state.persons_smplx = [_make_smplx_person(1)] + state.smplx_last_t = 1.0 + pub._tick(t_now=0.0) + bridge.reset_mock() + state.persons_smplx = [] + state.smplx_last_t = 2.0 + state.persons_body3d = [] + pub._tick(t_now=1.0) + bridge.send_leave.assert_called_with(pid=1) + + +def test_publisher_no_double_emit_same_timestamp() -> None: + from data_only_viz.action_head_pub import ActionHeadPublisher + state = _FakeState() + bridge = MagicMock() + pub = ActionHeadPublisher(state, bridge, ckpt_path=None) + state.persons_smplx = [_make_smplx_person(1)] + state.smplx_last_t = 1.0 + pub._tick(t_now=0.0) + bridge.reset_mock() + pub._tick(t_now=1.0) # same smplx_last_t + bridge.send_action.assert_not_called() + + +def test_publisher_uses_face_lips_for_mouth_open() -> None: + """mouth_open from MediaPipe lip landmarks (idx 13 and 14) must be ~1.0.""" + from unittest.mock import patch + from data_only_viz.action_head_pub import ActionHeadPublisher, MEDIAPIPE_LIP_UPPER_INNER, MEDIAPIPE_LIP_LOWER_INNER + state = _FakeState() + bridge = MagicMock() + pub = ActionHeadPublisher(state, bridge, ckpt_path=None) + + # Build a fake face landmark list: at least 15 landmarks. + # idx 13 = upper inner (y=0), idx 14 = lower inner (y=1), rest zeros. + face_kps = [(0.0, 0.0, 0.0)] * 15 + face_kps[MEDIAPIPE_LIP_UPPER_INNER] = (0.0, 0.0, 0.0) + face_kps[MEDIAPIPE_LIP_LOWER_INNER] = (1.0, 0.0, 0.0) # 1m apart in x + state.persons_face = [face_kps] + state.persons_face_ids = [0] + + captured_mouth: list[float] = [] + original_step = pub.head.step + + def spy_step(pid, j3d, expr=None, mouth_open=0.0, hands_kp=None): + captured_mouth.append(mouth_open) + return original_step(pid, j3d, expr=expr, mouth_open=mouth_open, hands_kp=hands_kp) + + pub.head.step = spy_step # type: ignore[method-assign] + + state.persons_smplx = [_make_smplx_person(0)] + state.smplx_last_t = 1.0 + pub._tick(t_now=0.0) + + assert len(captured_mouth) == 1 + assert abs(captured_mouth[0] - 1.0) < 1e-5 + + +def test_publisher_passes_hands_kp_to_step() -> None: + """hands_kp of shape (42, 3) must be passed to head.step.""" + from data_only_viz.action_head_pub import ActionHeadPublisher + state = _FakeState() + bridge = MagicMock() + pub = ActionHeadPublisher(state, bridge, ckpt_path=None) + + # Two 21-kp hand arrays (left + right) for pid=0. + rng = np.random.default_rng(7) + left_kps = rng.normal(size=(21, 3)).astype(np.float32) + right_kps = rng.normal(size=(21, 3)).astype(np.float32) + # persons_hands flat list: [left, right], ids both 0 (same pid). + state.persons_hands = [left_kps, right_kps] + state.persons_hands_ids = [0, 0] + + captured_hands: list = [] + original_step = pub.head.step + + def spy_step(pid, j3d, expr=None, mouth_open=0.0, hands_kp=None): + captured_hands.append(hands_kp) + return original_step(pid, j3d, expr=expr, mouth_open=mouth_open, hands_kp=hands_kp) + + pub.head.step = spy_step # type: ignore[method-assign] + + state.persons_smplx = [_make_smplx_person(0)] + state.smplx_last_t = 1.0 + pub._tick(t_now=0.0) + + assert len(captured_hands) == 1 + assert captured_hands[0] is not None + assert captured_hands[0].shape == (42, 3) diff --git a/data_only_viz/tests/test_augment.py b/data_only_viz/tests/test_augment.py new file mode 100644 index 0000000..c6b9a5a --- /dev/null +++ b/data_only_viz/tests/test_augment.py @@ -0,0 +1,46 @@ +"""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, 32, 3)).astype(np.float32) + + +def test_mirror_swap_left_right_joints() -> None: + from data_only_viz.training.augment import mirror_x, MIRROR_MAP + x = _sample_stack(0) + y = mirror_x(x) + # 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: + 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/tests/test_autolabel.py b/data_only_viz/tests/test_autolabel.py new file mode 100644 index 0000000..a529b1e --- /dev/null +++ b/data_only_viz/tests/test_autolabel.py @@ -0,0 +1,85 @@ +"""Tests for rule-based auto-labeler.""" +from __future__ import annotations + +import numpy as np + +from data_only_viz.action_head import WINDOW_LEN + + +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((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] + f[5] = [0.1, 0.4, 0.3] + f[7] = [-0.1, 0.1, 0.3] + f[8] = [0.1, 0.1, 0.3] + frames.append(f) + return frames + + +def _static_standing(frame_count: int = WINDOW_LEN) -> list[np.ndarray]: + """Hip high, knees ~180°.""" + frames = [] + for _ in range(frame_count): + 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] + f[5] = [0.1, 0.5, 0.0] + f[7] = [-0.1, 0.1, 0.0] + f[8] = [0.1, 0.1, 0.0] + frames.append(f) + return frames + + +def _dancing(frame_count: int = WINDOW_LEN) -> list[np.ndarray]: + """Standing pose with high wrist velocity.""" + base = _static_standing(1)[0] + frames = [] + for t in range(frame_count): + f = base.copy() + phase = 2 * np.pi * t * 0.125 # 0.125 = 1/8, slower oscillation + f[20] = base[20] + np.array([np.sin(phase) * 0.5, np.cos(phase) * 0.5, 0]) + f[21] = base[21] + np.array( + [-np.sin(phase) * 0.5, np.cos(phase) * 0.5, 0] + ) + frames.append(f.astype(np.float32)) + return frames + + +def test_autolabel_static_standing_is_debout() -> None: + from data_only_viz.training.autolabel import autolabel_window + + label, conf = autolabel_window(_static_standing()) + assert label == "debout" + assert conf >= 0.5 + + +def test_autolabel_static_seated_is_assise() -> None: + from data_only_viz.training.autolabel import autolabel_window + + label, conf = autolabel_window(_static_seated()) + assert label == "assise" + assert conf >= 0.5 + + +def test_autolabel_dancing_is_danse() -> None: + from data_only_viz.training.autolabel import autolabel_window + + label, conf = autolabel_window(_dancing()) + assert label == "danse" + assert conf >= 0.5 + + +def test_autolabel_ambiguous_is_none() -> None: + from data_only_viz.training.autolabel import autolabel_window + + base = _static_standing(WINDOW_LEN) + for t, f in enumerate(base): + f[20, 0] += 0.01 * np.sin(t) + label, _conf = autolabel_window(base) + assert label in ("debout", None) diff --git a/data_only_viz/tests/test_dataset.py b/data_only_viz/tests/test_dataset.py index 679cf66..64f3e08 100644 --- a/data_only_viz/tests/test_dataset.py +++ b/data_only_viz/tests/test_dataset.py @@ -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,10 +68,57 @@ 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) +def test_write_and_load_dataset_jsonl_with_hands_kp(tmp_path: Path) -> None: + from data_only_viz.training.dataset import ( + DatasetRow, + load_dataset_jsonl, + write_dataset_jsonl, + ) + rng = np.random.default_rng(1) + hands_kp = rng.normal(size=(16, 42, 3)).astype(np.float32) + row = DatasetRow( + window_id="sess01_pid1_w0000", + label="danse", + j3d_stack=rng.normal(size=(16, 32, 3)).astype(np.float32), + session="sess01", + pid_local=1, + auto_label_confidence=0.9, + manually_validated=True, + hands_kp_stack=hands_kp, + ) + out = tmp_path / "with_hands.jsonl" + write_dataset_jsonl([row], out) + loaded = load_dataset_jsonl(out) + assert loaded[0].hands_kp_stack is not None + assert loaded[0].hands_kp_stack.shape == (16, 42, 3) + assert np.allclose(loaded[0].hands_kp_stack, hands_kp, atol=1e-6) + + +def test_load_dataset_jsonl_without_hands_kp_is_ok(tmp_path: Path) -> None: + """Legacy v2 rows without hands_kp field should load with hands_kp_stack=None.""" + import json + from data_only_viz.training.dataset import load_dataset_jsonl + rng = np.random.default_rng(2) + row = { + "window_id": "sess01_pid1_w0000", + "label": "debout", + "j3d": rng.normal(size=(16, 32, 3)).tolist(), + "session": "sess01", + "pid_local": 1, + "auto_label_confidence": 0.8, + "manually_validated": False, + } + out = tmp_path / "legacy.jsonl" + out.write_text(json.dumps(row) + "\n") + loaded = load_dataset_jsonl(out) + assert len(loaded) == 1 + assert loaded[0].hands_kp_stack is None + + def test_split_by_session(tmp_path: Path) -> None: from data_only_viz.training.dataset import DatasetRow, split_by_session rng = np.random.default_rng(0) @@ -79,7 +126,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, )) diff --git a/data_only_viz/tests/test_extract_mediapipe_offline.py b/data_only_viz/tests/test_extract_mediapipe_offline.py new file mode 100644 index 0000000..0ba6e38 --- /dev/null +++ b/data_only_viz/tests/test_extract_mediapipe_offline.py @@ -0,0 +1,82 @@ +"""Sanity tests for MediaPipe offline extractor (no MediaPipe runtime -- we +mock the landmarker and feed synthetic landmarks).""" +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest + + +def test_build_j3d32_combines_body_and_fingertips() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _build_j3d32 + from data_only_viz.action_head import J3D_JOINTS + body3d = np.linspace(0, 1, 33 * 3).reshape(33, 3).astype(np.float32) + hands_kp42 = np.linspace(2, 3, 42 * 3).reshape(42, 3).astype(np.float32) + j3d = _build_j3d32(body3d, hands_kp42) + assert j3d is not None + assert j3d.shape == (J3D_JOINTS, 3) + # The body22 portion comes from body3d via MEDIAPIPE_TO_22. + # The fingertip portion (indices 22..31) comes from hands_kp at idx 4,8,12,16,20. + assert np.allclose(j3d[22], hands_kp42[4]) + assert np.allclose(j3d[26], hands_kp42[20]) + assert np.allclose(j3d[27], hands_kp42[21 + 4]) + assert np.allclose(j3d[31], hands_kp42[21 + 20]) + + +def test_build_j3d32_returns_none_when_no_body() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _build_j3d32 + j3d = _build_j3d32(None, np.zeros((42, 3), dtype=np.float32)) + assert j3d is None + + +def test_hands_kp42_combines_left_right_sides() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _hands_kp42 + left = np.linspace(0, 1, 21 * 3).reshape(21, 3).astype(np.float32) + right = np.linspace(2, 3, 21 * 3).reshape(21, 3).astype(np.float32) + out = _hands_kp42(left, right) + assert out.shape == (42, 3) + assert np.allclose(out[:21], left) + assert np.allclose(out[21:], right) + + +def test_hands_kp42_zero_pads_when_missing() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _hands_kp42 + left = np.ones((21, 3), dtype=np.float32) + out = _hands_kp42(left, None) + assert np.allclose(out[:21], left) + assert np.allclose(out[21:], 0.0) + + +def test_mouth_open_from_face_lips() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _mouth_open + # MediaPipe FaceMesh has 478 landmarks. Build a sparse array : zero + # everywhere except idx 13 (upper inner) and idx 14 (lower inner), + # 1 metre apart on the y axis. + face = np.zeros((478, 3), dtype=np.float32) + face[13] = [0.0, 1.0, 0.0] + face[14] = [0.0, 0.0, 0.0] + assert abs(_mouth_open(face) - 1.0) < 1e-6 + + +def test_mouth_open_returns_zero_on_empty_face() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _mouth_open + assert _mouth_open(None) == 0.0 + assert _mouth_open(np.zeros((10, 3), dtype=np.float32)) == 0.0 + + +def test_lmk_list_to_array_round_trip() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _lmk_list_to_array + class _Lmk: + def __init__(self, x: float, y: float, z: float) -> None: + self.x = x; self.y = y; self.z = z + lmks = [_Lmk(i, 2 * i, 3 * i) for i in range(5)] + arr = _lmk_list_to_array(lmks) + assert arr is not None + assert arr.shape == (5, 3) + assert np.allclose(arr[2], [2.0, 4.0, 6.0]) + + +def test_lmk_list_to_array_none_input() -> None: + from data_only_viz.scripts.extract_mediapipe_offline import _lmk_list_to_array + assert _lmk_list_to_array(None) is None diff --git a/data_only_viz/tests/test_training_smoke.py b/data_only_viz/tests/test_training_smoke.py new file mode 100644 index 0000000..7b9c74c --- /dev/null +++ b/data_only_viz/tests/test_training_smoke.py @@ -0,0 +1,63 @@ +"""Smoke test for action-head training (2 epochs, tiny dataset, CPU).""" +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + + +def _make_tiny_dataset(tmp_path: Path) -> Path: + from data_only_viz.training.dataset import DatasetRow, write_dataset_jsonl + rng = np.random.default_rng(0) + rows = [] + for sess_i, sess in enumerate(("s01", "s02", "s03")): + for w in range(30): + label = ("debout", "assise", "danse")[w % 3] + rows.append(DatasetRow( + window_id=f"{sess}_w{w:03d}", + label=label, + 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) + return out + + +def test_train_2_epochs_no_crash(tmp_path: Path) -> None: + from data_only_viz.training.train_action_head import train + ds = _make_tiny_dataset(tmp_path) + ckpt = tmp_path / "ckpt.pt" + history = train( + dataset_path=ds, + ckpt_out=ckpt, + epochs=2, + batch_size=8, + lr=1e-3, + device="cpu", + seed=0, + log_every=10_000, + ) + assert ckpt.exists() + assert len(history["train_loss"]) == 2 + assert all(np.isfinite(history["train_loss"])) + + +def test_trained_checkpoint_loadable(tmp_path: Path) -> None: + from data_only_viz.action_head import ActionHead + from data_only_viz.training.train_action_head import train + ds = _make_tiny_dataset(tmp_path) + ckpt = tmp_path / "ckpt.pt" + train(dataset_path=ds, ckpt_out=ckpt, epochs=1, batch_size=8, + 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((32, 3), dtype=np.float32)) + assert abs(float(probs.sum()) - 1.0) < 1e-5 diff --git a/data_only_viz/training/augment.py b/data_only_viz/training/augment.py new file mode 100644 index 0000000..0c3cff0 --- /dev/null +++ b/data_only_viz/training/augment.py @@ -0,0 +1,85 @@ +"""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 diff --git a/data_only_viz/training/autolabel.py b/data_only_viz/training/autolabel.py new file mode 100644 index 0000000..bbf63de --- /dev/null +++ b/data_only_viz/training/autolabel.py @@ -0,0 +1,120 @@ +"""Rule-based labeler for j3d windows. + +Outputs one of {"debout", "assise", "danse", None}. None marks +ambiguous windows that should be reviewed manually. + +Rules are tuned for SMPL-X joint indexing as used by Multi-HMR. +""" +from __future__ import annotations + +import argparse +import logging +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from data_only_viz.action_head import ( + FeatureExtractor, + HIP_LEFT, + HIP_RIGHT, + WINDOW_LEN, +) + + +@dataclass(frozen=True) +class AutoLabelConfig: + hip_y_seated_max: float = 0.55 + knee_angle_seated_max: float = 2.0 # rad, ~115° + speed_static_max: float = 0.03 # m/s mean joint speed + speed_dance_min: float = 0.033 # m/s mean joint speed + accel_dance_min: float = 0.001 + + +DEFAULT_CFG = AutoLabelConfig() + + +def autolabel_window( + frames: list[np.ndarray], cfg: AutoLabelConfig = DEFAULT_CFG +) -> tuple[str | None, float]: + """Return (label, confidence). label is None when ambiguous.""" + if len(frames) < WINDOW_LEN // 2: + return None, 0.0 + cur = frames[-1] + hip_y = float((cur[HIP_LEFT, 1] + cur[HIP_RIGHT, 1]) * 0.5) + knee_angle = FeatureExtractor._mean_knee_angle(cur) + kin = FeatureExtractor.kinetics(frames) + speed = float(kin[0]) + accel = float(kin[1]) + + if hip_y < cfg.hip_y_seated_max and knee_angle < cfg.knee_angle_seated_max: + conf = 0.5 + 0.5 * min(1.0, (cfg.hip_y_seated_max - hip_y) / 0.2) + return "assise", conf + if speed >= cfg.speed_dance_min or accel >= cfg.accel_dance_min: + conf = 0.5 + 0.5 * min(1.0, speed / 0.5) + return "danse", conf + if speed <= cfg.speed_static_max: + conf = 0.6 + return "debout", conf + return None, 0.0 + + +def autolabel_dataset( + frames_jsonl: Path, + out_jsonl: Path, + window_len: int = WINDOW_LEN, + stride: int = 4, + keep_none: bool = True, +) -> int: + """Glue: raw frames jsonl → sliding windows → auto-label → DatasetRow jsonl. + + Returns the number of windows written. + """ + from data_only_viz.training.dataset import ( + DatasetRow, + load_frames_jsonl, + sliding_windows, + write_dataset_jsonl, + ) + + frames = load_frames_jsonl(frames_jsonl) + rows = [] + for win in sliding_windows(frames, window_len=window_len, stride=stride): + frame_list = [win.j3d_stack[t] for t in range(win.j3d_stack.shape[0])] + label, conf = autolabel_window(frame_list) + if label is None and not keep_none: + continue + rows.append( + DatasetRow( + window_id=f"{win.session}_pid{win.pid_local}_t{int(win.first_ts*1000):08d}", + label=label if label is not None else "debout", + j3d_stack=win.j3d_stack, + session=win.session, + pid_local=win.pid_local, + auto_label_confidence=conf, + manually_validated=False, + ) + ) + out_jsonl.parent.mkdir(parents=True, exist_ok=True) + write_dataset_jsonl(rows, out_jsonl) + return len(rows) + + +def _cli() -> None: + p = argparse.ArgumentParser() + p.add_argument( + "--frames", + required=True, + type=Path, + help="Raw frames jsonl from extract_j3d_offline.py", + ) + p.add_argument("--out", required=True, type=Path, help="Auto-labeled windowed dataset jsonl") + p.add_argument("--stride", type=int, default=4) + args = p.parse_args() + logging.basicConfig(level=logging.INFO) + n = autolabel_dataset(args.frames, args.out, stride=args.stride) + print(f"wrote {n} windows to {args.out}") + + +if __name__ == "__main__": + _cli() diff --git a/data_only_viz/training/dataset.py b/data_only_viz/training/dataset.py index 60cd2cd..27f2ecd 100644 --- a/data_only_viz/training/dataset.py +++ b/data_only_viz/training/dataset.py @@ -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,35 @@ class RawFrame: ts: float session: str pid: int - j3d: np.ndarray # (22, 3) float32 + j3d: np.ndarray # (32, 3) float32 (v3: body22 + 10 fingertips) + expression: np.ndarray | None = None # (EXPR_DIM,) or None + mouth_open: float = 0.0 + hands_kp: np.ndarray | None = None # (42, 3) or None @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 + hands_kp_stack: np.ndarray | None = None # (window_len, 42, 3) 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 + hands_kp_stack: np.ndarray | None = None # (window_len, 42, 3) or None def load_frames_jsonl(path: Path) -> list[RawFrame]: @@ -45,11 +54,18 @@ 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 + hands_raw = d.get("hands_kp") + hands_kp = np.asarray(hands_raw, dtype=np.float32) if hands_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)), + hands_kp=hands_kp, )) return rows @@ -68,14 +84,43 @@ 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 + ) + # hands_kp stack: (window_len, 42, 3) if any frame has hands_kp + if any(c.hands_kp is not None for c in chunk): + hands_kp_stack = np.zeros((window_len, 42, 3), dtype=np.float32) + for t, c in enumerate(chunk): + if c.hands_kp is not None: + hk = np.asarray(c.hands_kp, dtype=np.float32) + if hk.shape == (42, 3): + hands_kp_stack[t] = hk + else: + hands_kp_stack = None 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, + hands_kp_stack=hands_kp_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 +128,14 @@ 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() + if r.hands_kp_stack is not None: + d["hands_kp_stack"] = r.hands_kp_stack.astype(np.float32).tolist() + f.write(json.dumps(d) + "\n") def load_dataset_jsonl(path: Path) -> list[DatasetRow]: @@ -94,6 +146,12 @@ 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 + hands_raw = d.get("hands_kp_stack") + hands_kp = np.asarray(hands_raw, dtype=np.float32) if hands_raw is not None else None out.append(DatasetRow( window_id=d["window_id"], label=d["label"], @@ -102,6 +160,9 @@ 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, + hands_kp_stack=hands_kp, )) return out diff --git a/data_only_viz/training/eval.py b/data_only_viz/training/eval.py new file mode 100644 index 0000000..c711187 --- /dev/null +++ b/data_only_viz/training/eval.py @@ -0,0 +1,87 @@ +"""Evaluate a trained action-head checkpoint.""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import DataLoader + +from data_only_viz.action_head import ActionHeadModel, LABELS +from data_only_viz.training.dataset import load_dataset_jsonl, split_by_session +from data_only_viz.training.train_action_head import ( + LABEL_TO_IDX, + WindowDataset, +) + + +def confusion_matrix(true: list[int], pred: list[int], + num_classes: int = 3) -> np.ndarray: + cm = np.zeros((num_classes, num_classes), dtype=np.int64) + for t, p in zip(true, pred): + cm[t, p] += 1 + return cm + + +def evaluate(ckpt_path: Path, dataset_path: Path, device: str = "cpu", + seed: int = 0) -> dict: + rows = load_dataset_jsonl(dataset_path) + _train, _val, test_rows = split_by_session(rows, seed=seed) + if not test_rows: + test_rows = rows + ds = WindowDataset(test_rows, augment=False, seed=seed) + loader = DataLoader(ds, batch_size=64, shuffle=False) + model = ActionHeadModel().to(device).eval() + payload = torch.load(ckpt_path, map_location=device, weights_only=True) + model.load_state_dict(payload["model_state_dict"]) + true: list[int] = [] + pred: list[int] = [] + with torch.no_grad(): + for x, y in loader: + x = x.to(device); y = y.to(device) + B, T, _ = x.shape + h = model.init_hidden(batch=B, device=device) + logits = None + for t in range(T): + logits, h = model(x[:, t, :], h) + true.extend(y.cpu().tolist()) + pred.extend(logits.argmax(-1).cpu().tolist()) + cm = confusion_matrix(true, pred) + acc = float(np.trace(cm) / max(1, cm.sum())) + confusion_db = float((cm[0, 2] + cm[2, 0]) / max(1, cm.sum())) + feat_dim = ds[0][0].shape[-1] + bench_x = torch.zeros(1, feat_dim, device=device) + h = model.init_hidden(batch=1, device=device) + for _ in range(20): + _ = model(bench_x, h) + t0 = time.perf_counter() + N = 500 + for _ in range(N): + _, h = model(bench_x, h) + lat_ms = (time.perf_counter() - t0) * 1000.0 / N + return { + "test_acc": acc, + "confusion_debout_danse": confusion_db, + "confusion_matrix": cm.tolist(), + "labels": list(LABELS), + "step_latency_ms": lat_ms, + "n_test": int(cm.sum()), + } + + +def _cli() -> None: + p = argparse.ArgumentParser() + p.add_argument("--ckpt", required=True, type=Path) + p.add_argument("--dataset", required=True, type=Path) + p.add_argument("--device", default="cpu", + choices=["cpu", "mps", "cuda"]) + args = p.parse_args() + out = evaluate(args.ckpt, args.dataset, device=args.device) + print(json.dumps(out, indent=2)) + + +if __name__ == "__main__": + _cli() diff --git a/data_only_viz/training/review.py b/data_only_viz/training/review.py new file mode 100644 index 0000000..2cbb848 --- /dev/null +++ b/data_only_viz/training/review.py @@ -0,0 +1,116 @@ +"""Manual label review TUI. + +Reads an auto-labeled jsonl dataset, presents each window with: + - ASCII skeleton (front view) of last frame + - speed/accel/sym kinetics + - proposed label + confidence +Keys: + 1 = debout, 2 = assise, 3 = danse + ENTER = accept proposed label + S = skip (label = None, will not be saved) + Q = quit and write what we have so far + +Usage: + uv run python -m data_only_viz.training.review \\ + --in ~/.cache/av-live-action/dataset/auto.jsonl \\ + --out ~/.cache/av-live-action/dataset/reviewed.jsonl +""" +from __future__ import annotations + +import argparse +import sys +import termios +import tty +from pathlib import Path + +import numpy as np + +from data_only_viz.action_head import LABELS +from data_only_viz.training.autolabel import autolabel_window +from data_only_viz.training.dataset import ( + DatasetRow, + load_dataset_jsonl, + write_dataset_jsonl, +) + + +def _ascii_skeleton(j3d: np.ndarray, width: int = 40, height: int = 16) -> str: + pts = j3d[:, [0, 1]] # x, y + mn = pts.min(axis=0) + mx = pts.max(axis=0) + rng = np.maximum(mx - mn, 1e-3) + norm = (pts - mn) / rng + grid = [[" "] * width for _ in range(height)] + for x, y in norm: + col = int(x * (width - 1)) + row = int((1 - y) * (height - 1)) + grid[row][col] = "*" + return "\n".join("".join(row) for row in grid) + + +def _getch() -> str: + fd = sys.stdin.fileno() + old = termios.tcgetattr(fd) + try: + tty.setraw(fd) + return sys.stdin.read(1) + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old) + + +def review(in_path: Path, out_path: Path, + sample_validated_fraction: float = 0.2, + seed: int = 0) -> None: + from data_only_viz.action_head import FeatureExtractor + + rows = load_dataset_jsonl(in_path) + rng = np.random.default_rng(seed) + kept: list[DatasetRow] = [] + for i, r in enumerate(rows): + proposed, conf = autolabel_window(list(r.j3d_stack)) + is_none = proposed is None + sampled = rng.random() < sample_validated_fraction + if not is_none and not sampled and r.manually_validated: + kept.append(r) + continue + print("\033[2J\033[H") # clear + print(f"[{i + 1}/{len(rows)}] {r.window_id} proposed={proposed} conf={conf:.2f}") + print(_ascii_skeleton(r.j3d_stack[-1])) + kin = FeatureExtractor.kinetics(list(r.j3d_stack)) + print(f"speed={kin[0]:.3f} accel={kin[1]:.3f} sym={kin[2]:+.3f}") + print("keys: 1=debout 2=assise 3=danse ENTER=accept S=skip Q=quit") + k = _getch().lower() + if k == "q": + break + if k == "s": + continue + if k == "\r": + chosen = proposed + elif k in ("1", "2", "3"): + chosen = LABELS[int(k) - 1] + else: + continue + if chosen is None: + continue + kept.append(DatasetRow( + window_id=r.window_id, label=chosen, j3d_stack=r.j3d_stack, + session=r.session, pid_local=r.pid_local, + auto_label_confidence=conf, manually_validated=True, + )) + out_path.parent.mkdir(parents=True, exist_ok=True) + write_dataset_jsonl(kept, out_path) + print(f"\nwrote {len(kept)} rows to {out_path}") + + +def _cli() -> None: + p = argparse.ArgumentParser() + p.add_argument("--in", dest="in_path", required=True, type=Path) + p.add_argument("--out", dest="out_path", required=True, type=Path) + p.add_argument("--sample-fraction", type=float, default=0.2) + args = p.parse_args() + review(args.in_path, args.out_path, + sample_validated_fraction=args.sample_fraction) + + +if __name__ == "__main__": + _cli() diff --git a/data_only_viz/training/train_action_head.py b/data_only_viz/training/train_action_head.py new file mode 100644 index 0000000..fd4ea53 --- /dev/null +++ b/data_only_viz/training/train_action_head.py @@ -0,0 +1,205 @@ +"""Train ActionHead on the windowed dataset. + +Usage: + uv run python -m data_only_viz.training.train_action_head \ + --dataset ~/.cache/av-live-action/dataset/dataset.jsonl \ + --ckpt-out ~/.cache/av-live-action/checkpoints/action_head.pt \ + --device mps --epochs 50 --batch-size 128 +""" +from __future__ import annotations + +import argparse +import json +import logging +from collections import Counter +from pathlib import Path + +import numpy as np +import torch +from torch import nn +from torch.utils.data import DataLoader, Dataset + +from data_only_viz.action_head import ( + ActionHeadModel, + EXPR_DIM, + FeatureExtractor, + HANDS_KP_FLAT, + HIP_LEFT, + HIP_RIGHT, + LABELS, +) +from data_only_viz.training.augment import random_augment +from data_only_viz.training.dataset import ( + DatasetRow, + load_dataset_jsonl, + split_by_session, +) + +LOG = logging.getLogger("train_action_head") +LABEL_TO_IDX = {l: i for i, l in enumerate(LABELS)} + + +class WindowDataset(Dataset[tuple[torch.Tensor, int]]): + def __init__(self, rows: list[DatasetRow], + augment: bool = False, seed: int = 0) -> None: + self._rows = rows + self._augment = augment + self._rng = np.random.default_rng(seed) + + def __len__(self) -> int: + return len(self._rows) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]: + row = self._rows[idx] + 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(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 + # hands_kp at frame t (42, 3); zeros if row has none + if row.hands_kp_stack is not None: + hands_t = row.hands_kp_stack[t] + hands_flat = hands_t.reshape(-1).astype(np.float32, copy=False) + else: + hands_flat = np.zeros(HANDS_KP_FLAT, dtype=np.float32) + feat = np.concatenate([ + cur.reshape(-1), vel.reshape(-1), accel.reshape(-1), + hands_flat, + 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 + prev = cur + x = torch.from_numpy(np.stack(feats)) + y = LABEL_TO_IDX[row.label] + return x, y + + +def _class_weights(rows: list[DatasetRow]) -> torch.Tensor: + counts = Counter(r.label for r in rows) + total = sum(counts.values()) + weights = torch.tensor([ + total / (len(LABELS) * counts.get(l, 1)) for l in LABELS + ], dtype=torch.float32) + return weights + + +def _run_epoch(model: nn.Module, loader: DataLoader, loss_fn: nn.Module, + optim: torch.optim.Optimizer | None, + device: str) -> tuple[float, float]: + train_mode = optim is not None + model.train(train_mode) + total_loss = 0.0 + correct = 0 + seen = 0 + for x, y in loader: + x = x.to(device) + y = y.to(device) + B, T, _ = x.shape + h = model.init_hidden(batch=B, device=device) + logits_last: torch.Tensor | None = None + for t in range(T): + logits, h = model(x[:, t, :], h) + logits_last = logits + assert logits_last is not None + loss = loss_fn(logits_last, y) + if train_mode: + optim.zero_grad() + loss.backward() + optim.step() + total_loss += float(loss) * B + correct += int((logits_last.argmax(-1) == y).sum()) + seen += B + return total_loss / max(1, seen), correct / max(1, seen) + + +def train(*, + dataset_path: Path, + ckpt_out: Path, + epochs: int = 50, + batch_size: int = 128, + lr: float = 1e-3, + device: str = "cpu", + seed: int = 0, + log_every: int = 1, + ) -> dict[str, list[float]]: + torch.manual_seed(seed) + rows = load_dataset_jsonl(dataset_path) + train_rows, val_rows, _test_rows = split_by_session(rows, seed=seed) + LOG.info("train=%d val=%d", len(train_rows), len(val_rows)) + train_ds = WindowDataset(train_rows, augment=True, seed=seed) + val_ds = WindowDataset(val_rows, augment=False, seed=seed) + train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True) + val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False) + model = ActionHeadModel().to(device) + weights = _class_weights(train_rows).to(device) + loss_fn = nn.CrossEntropyLoss(weight=weights) + optim = torch.optim.AdamW(model.parameters(), lr=lr) + history: dict[str, list[float]] = { + "train_loss": [], "train_acc": [], "val_loss": [], "val_acc": [], + } + best_val_acc = -1.0 + ckpt_out.parent.mkdir(parents=True, exist_ok=True) + for ep in range(epochs): + tl, ta = _run_epoch(model, train_loader, loss_fn, optim, device) + with torch.no_grad(): + vl, va = _run_epoch(model, val_loader, loss_fn, None, device) + history["train_loss"].append(tl) + history["train_acc"].append(ta) + history["val_loss"].append(vl) + history["val_acc"].append(va) + if ep % log_every == 0 or ep == epochs - 1: + LOG.info("ep=%d train_loss=%.4f train_acc=%.3f val_loss=%.4f val_acc=%.3f", + ep, tl, ta, vl, va) + if va > best_val_acc: + best_val_acc = va + torch.save({"model_state_dict": model.state_dict(), + "version": 1, "val_acc": va}, ckpt_out) + return history + + +def _cli() -> None: + p = argparse.ArgumentParser() + p.add_argument("--dataset", required=True, type=Path) + p.add_argument("--ckpt-out", required=True, type=Path) + p.add_argument("--epochs", type=int, default=50) + p.add_argument("--batch-size", type=int, default=128) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--device", default="cpu", + choices=["cpu", "mps", "cuda"]) + p.add_argument("--seed", type=int, default=0) + args = p.parse_args() + logging.basicConfig(level=logging.INFO, + format="%(asctime)s [%(name)s] %(message)s") + hist = train(dataset_path=args.dataset, ckpt_out=args.ckpt_out, + epochs=args.epochs, batch_size=args.batch_size, + lr=args.lr, device=args.device, seed=args.seed) + print(json.dumps({"final": {k: v[-1] for k, v in hist.items()}})) + + +if __name__ == "__main__": + _cli() diff --git a/docs/superpowers/plans/2026-05-13-action-head.md b/docs/superpowers/plans/2026-05-13-action-head.md index 389ce06..c06fa34 100644 --- a/docs/superpowers/plans/2026-05-13-action-head.md +++ b/docs/superpowers/plans/2026-05-13-action-head.md @@ -1,12 +1,18 @@ # action-head Implementation Plan > **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. +> +> **STATUS 2026-05-13 22:50** — Implementation **complete** (16/17 tasks, Task 16 is a manual gate). 39 tests green. Key deviations from this document, captured in the "Post-impl deviations" section below: +> - Task 14 pivoted from "modify `multi_hmr_worker_coreml.py` + CLI flag" to **standalone publisher thread `data_only_viz/action_head_pub.py`** + 3-line wire-in in `multi.py` (avoids collision with the user's parallel iteration on `multi_hmr_worker.py`). The MultiHMR backend is selected via env var `MULTIHMR_BACKEND=pytorch|coreml`, not a CLI flag. +> - Task 11 pivoted from "refactor `MultiHMRWorker` with `create_for_offline()`" to **standalone script using `MultiHMRCoreMLBackend.infer()` directly** — no worker refactor. +> - j3d is approximated from SMPL-X v3d via a fixed 22-vertex anchor set (`SMPLX_JOINT_ANCHOR_VERTS`), with a MediaPipe 33→22 fallback. The same anchor set is shared between live serve (`action_head_pub.py`) and offline extract (`scripts/extract_j3d_offline.py`) to avoid train/serve skew. +> - Studio train wrapper added as Task 8.5 (`data_only_viz/scripts/train_on_studio.sh`), validated end-to-end smoke 160 windows × 3 epochs MPS in ~4 s. **Goal:** Implement a real-time per-person action classifier (debout/assise/danse) on top of Multi-HMR `j3d`, with OSC output enriched by softmax probabilities and kinetics scalars (speed/accel/symmetry). **Architecture:** GRU-1-layer + MLP head streaming inference, fed by a 16-frame ring buffer per person. Trained windowed on Studio M3 Ultra (PyTorch MPS), inferred streaming on M5. Hybrid auto-labeler (rules on j3d) + manual review for dataset. Inference ≤ 2 ms/person M5 in eager PyTorch — no CoreML conversion needed. -**Tech Stack:** Python 3.11 + uv, PyTorch (MPS for train, CPU for M5 inference), numpy, python-osc, pytest. Reuses existing `data_only_viz` infrastructure (`multi_hmr_worker_coreml.py`, `pose_bridge.py`, `tracker.py`, `state.py`). +**Tech Stack:** Python 3.11 + uv, PyTorch (MPS for train, CPU for M5 inference), numpy, python-osc, pytest. Reuses existing `data_only_viz` infrastructure (`multi_hmr_worker.py`, `multihmr_coreml.py`, `pose_bridge.py`, `tracker.py`, `state.py`). **Reference spec:** `docs/superpowers/specs/2026-05-13-action-head-design.md` @@ -1839,6 +1845,8 @@ git commit -m "feat(data-only-viz): action capture script" ## Task 11 — Extract j3d offline +> **SUPERSEDED 2026-05-13.** The implemented script does NOT refactor `multi_hmr_worker.py`. Instead it uses the standalone `MultiHMRCoreMLBackend.infer()` from `data_only_viz/multihmr_coreml.py` directly. Output jsonl rows contain a (22, 3) `j3d` extracted via `SMPLX_JOINT_ANCHOR_VERTS` (shared with `action_head_pub.py` to avoid train/serve skew), not the raw v3d. See actual file at `data_only_viz/scripts/extract_j3d_offline.py`. The body below documents the original intent — keep as historical context. + **Files:** - Create: `data_only_viz/scripts/extract_j3d_offline.py` @@ -2196,6 +2204,8 @@ git commit -m "feat(data-only-viz): pose_bridge /pose/action + /pose/kin" ## Task 14 — Wire ActionHead into multi_hmr_worker_coreml +> **SUPERSEDED 2026-05-13.** No `multi_hmr_worker_coreml.py` file exists in the current repo — the user pivoted to `multihmr_coreml.py` (standalone backend) selected via env `MULTIHMR_BACKEND=pytorch|coreml` inside `multi_hmr_worker.py`. To avoid colliding with that file under active iteration, ActionHead wiring was implemented as a standalone publisher thread in `data_only_viz/action_head_pub.py` plus a 3-line wire-in inside `data_only_viz/multi.py` (`__init__` instantiates and `.start()`s the publisher). The publisher polls `state.persons_smplx` (preferred) and `state.persons_body3d` (MediaPipe fallback) at 30 Hz, deduplicates by timestamp, extracts j3d22 via shared `SMPLX_JOINT_ANCHOR_VERTS` / `MEDIAPIPE_TO_22` index maps, runs `ActionHead.step()` per pid, and emits OSC via the existing `PoseSoundBridge`. No CLI flag was added. See `data_only_viz/action_head_pub.py` and `data_only_viz/multi.py:22,97-98`. The body below documents the original intent — keep as historical context. + **Files:** - Modify: `data_only_viz/multi_hmr_worker_coreml.py` - Modify: `data_only_viz/main.py` diff --git a/docs/superpowers/specs/2026-05-13-action-head-design.md b/docs/superpowers/specs/2026-05-13-action-head-design.md index 7c262e2..e62a35e 100644 --- a/docs/superpowers/specs/2026-05-13-action-head-design.md +++ b/docs/superpowers/specs/2026-05-13-action-head-design.md @@ -1,11 +1,17 @@ # action-head — Classifier d'action temps réel au-dessus de Multi-HMR > **Date** : 2026-05-13 -> **Status** : design approuvé, prêt pour implementation plan +> **Status** : design approuvé — **implémenté 2026-05-13 22:50**, 16/17 tasks, 39 tests verts. Task 16 (E2E gate) reste manuel (requiert capture + train réel). > **Authors** : L'Electron Rare + Claude > **Companion plans** : > - `2026-05-13-multihmr-coreml-hybrid-backbone.md` > - `2026-05-13-studio-train-deploy-m5.md` +> +> **Déviations notables vs design original** (cf. plan `2026-05-13-action-head.md` pour le détail) : +> - **Wiring worker** : standalone publisher thread `data_only_viz/action_head_pub.py` + 3 lignes dans `multi.py`, au lieu de modifier directement `multi_hmr_worker.py` (qui était en cours d'évolution par l'utilisateur en parallèle). Backend Multi-HMR sélectionné par env `MULTIHMR_BACKEND=pytorch|coreml`, pas par flag CLI. +> - **Source j3d** : approximée via 22 vertex anchors (`SMPLX_JOINT_ANCHOR_VERTS`) sur le mesh SMPL-X 10475-vert, partagés entre serve live (`action_head_pub.py`) et extraction offline (`scripts/extract_j3d_offline.py`) pour éviter le train/serve skew. Fallback MediaPipe 33→22 (`MEDIAPIPE_TO_22`) quand `persons_smplx` est vide. **Limitation** : ces 22 indices sont approximatifs ; pour des j3d SMPL-X corrects, brancher `J_regressor @ v3d` quand le module SMPL-X est dispo. +> - **Extract offline** : pas de refactor de `MultiHMRWorker`, on utilise `MultiHMRCoreMLBackend.infer()` directement (commit user `9e7a9f8`). +> - **Studio launch** : wrapper bash `data_only_viz/scripts/train_on_studio.sh` (Task 8.5) qui rsync + ssh + uv sync + train MPS + ckpt back. Validé end-to-end sur dataset smoke 160 windows × 3 epochs en ~4 s wallclock. ## TL;DR @@ -94,7 +100,9 @@ class ActionHead: def forget(self, pid: int) -> None: ... ``` -Aucune modification de l'API publique de `multi_hmr_worker_coreml.py` autre que : +**Note d'implémentation 2026-05-13** : la section ci-dessous décrit l'intention originale. L'implémentation réelle est dans `data_only_viz/action_head_pub.py` (publisher thread) — pas de modification de `multi_hmr_worker.py`. Voir l'en-tête du document pour les déviations. + +Aucune modification de l'API publique de `multi_hmr_worker.py` n'est requise au-delà de : - Construction d'une `ActionHead` au startup. - Appel `.step()` après chaque détection. - Appel `.forget()` synchronisé avec `tracker.purge()`. diff --git a/sound_algo/data_only/scene_pose_action.scd b/sound_algo/data_only/scene_pose_action.scd new file mode 100644 index 0000000..096c119 --- /dev/null +++ b/sound_algo/data_only/scene_pose_action.scd @@ -0,0 +1,131 @@ +// ===================================================================== +// scene_pose_action.scd -- pilote les FX live a partir de l'action_head. +// +// Necessite : data_feeds.scd deja charge (handlers OSCdef \poseAction, +// \poseKin, \poseEnter, \poseLeave). Les dicts ~poseState et ~poseKin +// doivent etre populates en temps reel par data_only_viz.action_head_pub. +// +// Mapping : +// speed (m/s mean joint) -> drive amount 0..1 +// accel (m/s2) -> filter cutoff 200 Hz..6 kHz +// symmetry (-1..1) -> stereo width 0..1 +// label argmax (0=debout, 1=assise, 2=danse) : +// - debout -> reverb mid, kick on, melody pad +// - assise -> pad ambient, kick off, lo-fi +// - danse -> drive max, kick punchy, acid lead +// +// Usage live (a evaluer bloc par bloc dans le SC IDE) : +// [0] SETUP : declare le helper ~mapPoseToFx + lance la Routine. +// [1] START : ~scenePoseAction.start +// [2] STOP : ~scenePoseAction.stop +// [3] TWEAK : ajuster les seuils dans ~paConfig +// ===================================================================== + +// [0] SETUP ----------------------------------------------------------- +( +~paConfig = ( + speedMin: 0.0, speedMax: 0.8, + accelMin: 0.0, accelMax: 3.0, + cutoffMin: 200, cutoffMax: 6000, + rate: 10, // refresh Hz +); + +~mapPoseToFx = { + var avgSpeed = 0, avgAccel = 0, avgSym = 0, avgProbs = [0, 0, 0]; + var n = max(1, ~poseKin.size); + var labelIdx; + + ~poseKin.values.do { |kin| + avgSpeed = avgSpeed + (kin[\speed] ? 0); + avgAccel = avgAccel + (kin[\accel] ? 0); + avgSym = avgSym + (kin[\symmetry] ? 0); + }; + avgSpeed = avgSpeed / n; + avgAccel = avgAccel / n; + avgSym = avgSym / n; + + ~poseState.values.do { |ps| + var p = ps[\probs] ? [0, 0, 0]; + avgProbs = avgProbs.collect { |v, i| v + (p[i] ? 0) }; + }; + avgProbs = avgProbs.collect { _ / n }; + labelIdx = avgProbs.maxIndex ? 0; + + // Drive + filter + width + if (~fxDrive.notNil) { + ~fxDrive.(avgSpeed.linlin(~paConfig[\speedMin], ~paConfig[\speedMax], 0, 1)); + }; + if (~fxCut.notNil) { + ~fxCut.(avgAccel.linexp(~paConfig[\accelMin], ~paConfig[\accelMax], + ~paConfig[\cutoffMin], ~paConfig[\cutoffMax])); + }; + if (~fxSt.notNil) { ~fxSt.(avgSym.linlin(-1, 1, 0, 1)) }; + + // Label-driven track switch (smooth crossfade between presets) + switch (labelIdx, + 0, { + if (~fxComp.notNil) { ~fxComp.(0.7) }; + if (~fxRev.notNil) { ~fxRev.(0.4) }; + }, + 1, { + if (~fxComp.notNil) { ~fxComp.(0.3) }; + if (~fxRev.notNil) { ~fxRev.(0.6) }; + }, + 2, { + if (~fxComp.notNil) { ~fxComp.(0.9) }; + if (~fxRev.notNil) { ~fxRev.(0.2) }; + if (~fxDrive.notNil){ ~fxDrive.(avgSpeed.linlin(0, 0.6, 0.3, 1.0)) }; + } + ); + + [labelIdx, avgSpeed.round(0.01), avgAccel.round(0.01), + avgSym.round(0.01)]; +}; + +~scenePoseAction = ~scenePoseAction ?? { + ( + running: false, + routine: nil, + start: { + if (~scenePoseAction[\running]) { + "[scene_pose_action] already running".postln; + } { + ~scenePoseAction[\running] = true; + ~scenePoseAction[\routine] = Routine({ + inf.do { + var snapshot = ~mapPoseToFx.(); + if (~scenePoseAction[\verbose] ? false) { + ("[pose] label=" ++ snapshot[0] + ++ " s=" ++ snapshot[1] + ++ " a=" ++ snapshot[2] + ++ " sym=" ++ snapshot[3]).postln; + }; + (1 / ~paConfig[\rate]).wait; + }; + }).play; + "[scene_pose_action] STARTED".postln; + }; + }, + stop: { + ~scenePoseAction[\routine] !? { |r| r.stop }; + ~scenePoseAction[\routine] = nil; + ~scenePoseAction[\running] = false; + "[scene_pose_action] STOPPED".postln; + }, + verbose: false, + ) +}; + +"[OK] SETUP scene_pose_action".postln; +) + +// [1] START ----------------------------------------------------------- +// ~scenePoseAction[\verbose] = true; // optionnel : log chaque tick +// ~scenePoseAction[\start].(); + +// [2] STOP ------------------------------------------------------------ +// ~scenePoseAction[\stop].(); + +// [3] TWEAK ------------------------------------------------------------ +// ~paConfig[\speedMax] = 1.2; +// ~paConfig[\rate] = 20;