feat(data-only-viz): action auto-labeler rules
Problem: Action classification (debout/assise/danse) requires rule-based labeling before neural training. Task 5 of action-head plan. Approach: Heuristic rules on j3d posture + kinetics (speed/accel): - Hip height + knee angle → seated vs. standing - Joint velocity → static vs. dancing - Confidence scoring for ambiguous windows Changes: - action_head.py: scaffold with FeatureExtractor (kinetics, knee angle) - autolabel.py: AutoLabelConfig, autolabel_window(), autolabel_dataset() CLI glue (raw frames jsonl → windowed labeled dataset jsonl) - test_autolabel.py: 4 TDD tests (debout, assise, danse, ambiguous) Impact: Enables dataset creation pipeline (extract_j3d → auto-label → manual review → train ActionHead GRU).
This commit is contained in:
L'électron rare
@@ -0,0 +1,102 @@
|
||||
"""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, accel_mag, symmetry_score).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
# Constants (SMPL-X joint indexing as used by Multi-HMR)
|
||||
WINDOW_LEN: int = 16
|
||||
J3D_JOINTS: int = 22
|
||||
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
|
||||
|
||||
# Joint indices (SMPL-X)
|
||||
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
|
||||
|
||||
|
||||
class FeatureExtractor:
|
||||
"""Extract kinematic features from j3d window."""
|
||||
|
||||
@staticmethod
|
||||
def _mean_knee_angle(j3d: np.ndarray) -> float:
|
||||
"""Estimate mean knee angle (radians) from two frames.
|
||||
|
||||
j3d : (22, 3) float32
|
||||
Returns: angle in radians (0 = fully extended, π ≈ fully bent)
|
||||
"""
|
||||
hip_l = j3d[HIP_LEFT]
|
||||
knee_l = j3d[KNEE_LEFT]
|
||||
ankle_l = j3d[ANKLE_LEFT]
|
||||
|
||||
# Vectors: hip→knee, knee→ankle
|
||||
v1 = knee_l - hip_l
|
||||
v2 = ankle_l - knee_l
|
||||
|
||||
norm1 = np.linalg.norm(v1)
|
||||
norm2 = np.linalg.norm(v2)
|
||||
|
||||
if norm1 < 1e-6 or norm2 < 1e-6:
|
||||
return np.pi / 2 # neutral default
|
||||
|
||||
cos_angle = np.dot(v1, v2) / (norm1 * norm2)
|
||||
cos_angle = np.clip(cos_angle, -1.0, 1.0)
|
||||
angle = np.arccos(cos_angle)
|
||||
return float(angle)
|
||||
|
||||
@staticmethod
|
||||
def kinetics(frames: list[np.ndarray]) -> tuple[float, float, float]:
|
||||
"""Compute speed, accel, symmetry from frame window.
|
||||
|
||||
frames : list of (22, 3) float32 arrays
|
||||
Returns: (speed m/s, accel m/s², symmetry -1..1)
|
||||
"""
|
||||
if len(frames) < 2:
|
||||
return 0.0, 0.0, 0.0
|
||||
|
||||
# Speed: mean joint velocity magnitude
|
||||
velocities = []
|
||||
for i in range(1, len(frames)):
|
||||
dj3d = frames[i] - frames[i - 1]
|
||||
vel_mag = np.linalg.norm(dj3d, axis=1).mean()
|
||||
velocities.append(vel_mag)
|
||||
|
||||
speed = float(np.mean(velocities)) if velocities else 0.0
|
||||
|
||||
# Accel: finite difference of velocities
|
||||
accel = 0.0
|
||||
if len(velocities) >= 2:
|
||||
accels = np.abs(np.diff(velocities))
|
||||
accel = float(np.mean(accels)) if len(accels) > 0 else 0.0
|
||||
|
||||
# Symmetry: cosine similarity left/right shoulder and wrist
|
||||
cur = frames[-1]
|
||||
left_arm = np.concatenate([cur[SHOULDER_LEFT], cur[WRIST_LEFT]])
|
||||
right_arm = np.concatenate([cur[SHOULDER_RIGHT], cur[WRIST_RIGHT]])
|
||||
|
||||
norm_l = np.linalg.norm(left_arm)
|
||||
norm_r = np.linalg.norm(right_arm)
|
||||
symmetry = 0.0
|
||||
if norm_l > 1e-6 and norm_r > 1e-6:
|
||||
symmetry = float(np.dot(left_arm, right_arm) / (norm_l * norm_r))
|
||||
|
||||
return speed, accel, symmetry
|
||||
@@ -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((22, 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((22, 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)
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user