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

feat(data-only-viz): action-head v3 hands+lips

Browse Source 
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
This commit is contained in:
L'électron rare
2026-05-13 23:26:14 +02:00
parent 623c47983d
commit beb94d2a4c
9 changed files with 298 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data_only_viz/action_head.py
+31 -12
View File
@@ -32,13 +32,20 @@ LABELS: tuple[str, str, str] = ("debout", "assise", "danse")
EXPR_DIM: int = 10
EXTRA_SCALARS: int = 4 # hip_y, knee_angle, sym_score, mouth_open
# Layout per step:
# [0 : 96 ] j3d (32, 3)
# 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 : 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
# [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
@@ -60,18 +67,20 @@ FINGERTIP_RIGHT_BASE: int = 27
class FeatureExtractor:
"""Stateless feature builder over a list of recent j3d frames.
Vector layout (FEATURE_DIM = 302):
[0 : 96 ] j3d current frame, flattened (32 joints x 3 dims)
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 : 298] expression PCA coefficients (10,)
[298 : 302] kinetics scalars (hip_y, knee_angle, symmetry_score, mouth_open)
[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) -> np.ndarray:
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]
@@ -83,6 +92,13 @@ 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)
# 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:
@@ -93,6 +109,7 @@ class FeatureExtractor:
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)
@@ -227,7 +244,8 @@ class ActionHead:
def step(self, pid: int, j3d: np.ndarray,
expr: np.ndarray | None = None,
mouth_open: float = 0.0) -> tuple[str, np.ndarray, np.ndarray]:
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
@@ -241,7 +259,8 @@ 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, expr=expr, mouth_open=mouth_open)
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:
data_only_viz/action_head_pub.py
+110 -40
View File
@@ -17,6 +17,9 @@ 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,
@@ -28,12 +31,12 @@ DEFAULT_CKPT = (
Path.home() / ".cache" / "av-live-action" / "checkpoints" / "action_head.pt"
)
# 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.
# 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, ...] = (
7174, 7397, 7670, 7942, 8214, # L
4631, 4854, 5127, 5399, 5671, # R
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:
@@ -56,6 +59,11 @@ assert len(SMPLX_JOINT_ANCHOR_VERTS) == 32
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)
@@ -120,11 +128,11 @@ class ActionHeadPublisher(threading.Thread):
self._last_body_t = source_t
current_pids: set[int] = set()
if persons32:
for pid, j3d, expr, mouth in 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,
mouth_open=mouth)
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)
@@ -135,13 +143,13 @@ class ActionHeadPublisher(threading.Thread):
self.bridge.send_leave(pid=gone)
self._last_pids = current_pids
def _read_sources(
self,
) -> tuple[list[tuple[int, np.ndarray, np.ndarray, float]] | None,
float, str, bool]:
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).
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.
"""
@@ -150,12 +158,24 @@ class ActionHeadPublisher(threading.Thread):
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)
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.
# 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]] = []
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")
@@ -177,23 +197,24 @@ class ActionHeadPublisher(threading.Thread):
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_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
out.append((pid, j3d32, expr_np, mouth))
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 [])))
# 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
@@ -201,25 +222,74 @@ class ActionHeadPublisher(threading.Thread):
if arr is None or arr.shape[0] < 33:
continue
body22 = arr[list(MEDIAPIPE_TO_22)].astype(np.float32)
# fingertips from hands if available
# fingertips from persons_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]
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)
mouth = 0.0
out.append((pid, j3d32, expr_np, mouth))
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."""
data_only_viz/scripts/extract_j3d_offline.py
+4 -1
View File
@@ -17,7 +17,7 @@ from pathlib import Path
import cv2
import numpy as np
from data_only_viz.action_head import EXPR_DIM
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,
@@ -125,6 +125,9 @@ def extract(session: str, video: Path, out: Path,
"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
data_only_viz/tests/test_action_head_features.py
+4 -2
View File
@@ -12,7 +12,9 @@ def test_module_imports() -> None:
assert hasattr(action_head, "ActionHead")
assert action_head.WINDOW_LEN == 16
assert action_head.J3D_JOINTS == 32
assert action_head.FEATURE_DIM == 302
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")
@@ -67,7 +69,7 @@ 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 == (302,)
assert feat.shape == (428,)
assert feat.dtype == np.float32
assert not np.isnan(feat).any()
data_only_viz/tests/test_action_head_model.py
+2 -2
View File
@@ -24,11 +24,11 @@ def test_model_forward_shape() -> None:
assert h_new.shape == h.shape
def test_model_param_count_under_80k() -> None:
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 < 80_000, f"too many params: {n}"
assert n < 100_000, f"too many params: {n}"
def test_action_head_step_warmup_returns_debout() -> None:
data_only_viz/tests/test_action_head_pub.py
+68
View File
@@ -19,6 +19,8 @@ class _FakeState:
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):
@@ -84,3 +86,69 @@ def test_publisher_no_double_emit_same_timestamp() -> None:
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)
data_only_viz/tests/test_dataset.py
+47
View File
@@ -72,6 +72,53 @@ def test_write_and_load_dataset_jsonl(tmp_path: Path) -> None:
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)
data_only_viz/training/dataset.py
+24 -2
View File
@@ -15,9 +15,10 @@ class RawFrame:
ts: float
session: str
pid: int
j3d: np.ndarray # (32, 3) float32 (v2: body22 + 10 fingertips)
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
@@ -28,6 +29,7 @@ class WindowRow:
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
@@ -41,6 +43,7 @@ class DatasetRow:
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]:
@@ -53,6 +56,8 @@ def load_frames_jsonl(path: Path) -> list[RawFrame]:
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"]),
@@ -60,6 +65,7 @@ def load_frames_jsonl(path: Path) -> list[RawFrame]:
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
@@ -94,10 +100,21 @@ def sliding_windows(frames: list[RawFrame],
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,
expr_stack=expr_stack,
mouth_open_stack=mouth_stack)
mouth_open_stack=mouth_stack,
hands_kp_stack=hands_kp_stack)
def write_dataset_jsonl(rows: Iterable[DatasetRow], path: Path) -> None:
@@ -116,6 +133,8 @@ def write_dataset_jsonl(rows: Iterable[DatasetRow], path: Path) -> 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")
@@ -131,6 +150,8 @@ def load_dataset_jsonl(path: Path) -> list[DatasetRow]:
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"],
@@ -141,6 +162,7 @@ def load_dataset_jsonl(path: Path) -> list[DatasetRow]:
manually_validated=bool(d["manually_validated"]),
expr_stack=expr,
mouth_open_stack=mouth,
hands_kp_stack=hands_kp,
))
return out
data_only_viz/training/train_action_head.py
+8
View File
@@ -23,6 +23,7 @@ from data_only_viz.action_head import (
ActionHeadModel,
EXPR_DIM,
FeatureExtractor,
HANDS_KP_FLAT,
HIP_LEFT,
HIP_RIGHT,
LABELS,
@@ -78,8 +79,15 @@ class WindowDataset(Dataset[tuple[torch.Tensor, int]]):
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)