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