L'électron rare
beb94d2a4c
Extends the action-head feature pipeline from v2 (302-D) to v3 (428-D). - Replace placeholder SMPLX_FINGERTIP_VERTS with canonical vertex IDs from smplx.vertex_ids (lthumb/lindex/lmiddle/lring/lpinky, mirrored R) - Add HANDS_KP_* constants (21 kp/hand, 42 total, 126-D flat block) - FEATURE_DIM: 302 -> 428; hands_kp block inserted at [288:414] - FeatureExtractor.from_buffer gains hands_kp param (42, 3), zero-padded when absent - ActionHead.step gains hands_kp param, threads to from_buffer - _read_sources returns 5-tuples with hands_kp42x3 per person - MediaPipe FaceMesh inner-lip (idx 13/14) used for mouth_open; fallback to SMPL-X v3d lip vertices when face not available - _build_hands_map and _build_face_mouth_map helpers added - dataset.py: RawFrame/WindowRow/DatasetRow gain hands_kp fields - train_action_head.py: reads hands_kp_stack per step, zeros if absent - extract_j3d_offline.py: writes zero-filled hands_kp in jsonl output - Tests: FEATURE_DIM 302->428, param bound 80k->100k, +4 new tests
280 lines
10 KiB
Python
280 lines
10 KiB
Python
"""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)
|