cozyclay 1.10.0 → 2.0.0

package/tools/track/py/scene.py

@@ -0,0 +1,326 @@
+"""Scene terms for the known-character tracker: boxes, floor, stance, occlusion.
+
+World frame: Studio / three.js, metres, +Y up, floor plane y = floor_y (0).
+Boxes are oriented about +Y only (props stay upright, `src/ardy/collision-blockers.js`):
+a box point is `center + Ry(yaw) @ local` with `Ry` the three.js rotation
+(`tools/gt-render/scene-box.mjs` boxCorners), `|local| <= halfExtents`.
+
+Accepted box records (`parse_scene`):
+  - today's bench scene.json: `{min, max}` (plus optional `centre`, `halfExtents`, `yawDeg`)
+  - shot context: `{center|centre, halfExtents, yaw (radians) | yawDeg, perFrame?}` where
+    `perFrame` is a list of per-frame centres `[x, y, z]` or `{center|centre, yaw|yawDeg}`
+    for props that travel or turn.
+  - a list of such records, or `{boxes: [...]}`.
+
+Everything that feeds a loss takes plain torch tensors (vertices or joints) and is
+differentiable; the HMM and ray tests are discrete and run without gradients.
+"""
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import NamedTuple
+
+import numpy as np
+import torch
+
+# cskel27 indices: LeftFoot, LeftToeBase / RightFoot, RightToeBase (src/ardy/cskel27.js).
+FOOT_NAMES = ("left", "right")
+FOOT_JOINTS = ((25, 26), (21, 22))
+
+
+class SceneError(ValueError):
+    """Malformed scene input. `code` is the machine-readable reason."""
+
+    def __init__(self, code: str, detail: str):
+        super().__init__(f"{code}: {detail}")
+        self.code = code
+
+
+class Box(NamedTuple):
+    center: np.ndarray  # (3,)
+    half_extents: np.ndarray  # (3,)
+    yaw: float  # radians about +Y
+    per_frame_center: np.ndarray | None  # (T, 3)
+    per_frame_yaw: np.ndarray | None  # (T,)
+
+
+class SceneBoxes(NamedTuple):
+    centers: torch.Tensor  # (B, T, 3)
+    half: torch.Tensor  # (B, 3)
+    yaw: torch.Tensor  # (B, T)
+
+
+def _vec3(value, what: str) -> np.ndarray:
+    try:
+        out = np.asarray(value, dtype=np.float64)
+    except (TypeError, ValueError):
+        raise SceneError("bad-box", f"{what} must be 3 numbers, got {value!r}") from None
+    if out.shape != (3,) or not np.all(np.isfinite(out)):
+        raise SceneError("bad-box", f"{what} must be 3 finite numbers, got {value!r}")
+    return out
+
+
+def _yaw(record: dict, default: float, what: str) -> float:
+    if "yaw" in record:
+        value = record["yaw"]
+    elif "yawDeg" in record:
+        value = record["yawDeg"] * math.pi / 180 if isinstance(record["yawDeg"], (int, float)) else record["yawDeg"]
+    else:
+        return default
+    if not isinstance(value, (int, float)) or not math.isfinite(value):
+        raise SceneError("bad-box", f"{what} yaw must be a finite number, got {value!r}")
+    return float(value)
+
+
+def parse_box(record, index: int = 0) -> Box:
+    what = f"box {index}"
+    if not isinstance(record, dict):
+        raise SceneError("bad-box", f"{what} must be an object, got {type(record).__name__}")
+    center = record.get("center", record.get("centre"))
+    if center is not None or "halfExtents" in record:
+        if center is None or "halfExtents" not in record:
+            raise SceneError("bad-box", f"{what} needs both center and halfExtents")
+        center, half = _vec3(center, f"{what} center"), _vec3(record["halfExtents"], f"{what} halfExtents")
+    elif "min" in record and "max" in record:
+        lo, hi = _vec3(record["min"], f"{what} min"), _vec3(record["max"], f"{what} max")
+        center, half = (lo + hi) / 2, (hi - lo) / 2
+    else:
+        raise SceneError("bad-box", f"{what} needs {{center, halfExtents}} or {{min, max}}")
+    if not np.all(half > 0):
+        raise SceneError("degenerate-box", f"{what} has non-positive size {(2 * half).tolist()}")
+    yaw = _yaw(record, 0.0, what)
+    per_center = per_yaw = None
+    frames = record.get("perFrame")
+    if frames is not None:
+        if not isinstance(frames, list) or not frames:
+            raise SceneError("bad-box", f"{what} perFrame must be a non-empty list")
+        per_center, per_yaw = np.empty((len(frames), 3)), np.empty(len(frames))
+        for t, entry in enumerate(frames):
+            if isinstance(entry, dict):
+                per_center[t] = _vec3(entry.get("center", entry.get("centre")), f"{what} perFrame[{t}] center")
+                per_yaw[t] = _yaw(entry, yaw, f"{what} perFrame[{t}]")
+            else:
+                per_center[t], per_yaw[t] = _vec3(entry, f"{what} perFrame[{t}]"), yaw
+    return Box(center, half, yaw, per_center, per_yaw)
+
+
+def parse_scene(scene) -> list[Box]:
+    """Boxes from a scene.json / shot-context value (None -> no boxes)."""
+    if scene is None:
+        return []
+    if isinstance(scene, dict) and "boxes" in scene:
+        scene = scene["boxes"]
+    records = scene if isinstance(scene, list) else [scene]
+    return [parse_box(record, i) for i, record in enumerate(records)]
+
+
+def load_scene(path) -> list[Box]:
+    return parse_scene(json.loads(Path(path).read_text()))
+
+
+def box_tensors(boxes: list[Box], frames: int, *, device=None, dtype=torch.float32) -> SceneBoxes:
+    """Per-frame box tensors; static boxes are repeated over `frames`."""
+    centers, yaws = [], []
+    for i, box in enumerate(boxes):
+        if box.per_frame_center is not None:
+            if len(box.per_frame_center) != frames:
+                raise SceneError("box-frames-mismatch", f"box {i} has {len(box.per_frame_center)} perFrame entries for {frames} frames")
+            centers.append(box.per_frame_center)
+            yaws.append(box.per_frame_yaw)
+        else:
+            centers.append(np.broadcast_to(box.center, (frames, 3)))
+            yaws.append(np.full(frames, box.yaw))
+    as_t = lambda a, shape: torch.as_tensor(np.asarray(a, dtype=np.float64).reshape(shape), dtype=dtype, device=device)
+    b = len(boxes)
+    return SceneBoxes(
+        centers=as_t(centers, (b, frames, 3)),
+        half=as_t([box.half_extents for box in boxes], (b, 3)),
+        yaw=as_t(yaws, (b, frames)),
+    )
+
+
+def _boxes_like(boxes: SceneBoxes, points: torch.Tensor) -> SceneBoxes:
+    """Boxes on the points' device and dtype (differentiable; a no-op when they already match)."""
+    return SceneBoxes(*(v.to(device=points.device, dtype=points.dtype) for v in boxes))
+
+
+def _cpu64(value) -> torch.Tensor:
+    """Detached float64 CPU tensor from a tensor on any device, an array or a list."""
+    if isinstance(value, torch.Tensor):
+        return value.detach().to(device="cpu", dtype=torch.float64)
+    return torch.as_tensor(np.asarray(value, dtype=np.float64))
+
+
+def to_box_local(points: torch.Tensor, boxes: SceneBoxes) -> torch.Tensor:
+    """points (T, N, 3) -> box-local coordinates (T, N, B, 3), on the points' device."""
+    boxes = _boxes_like(boxes, points)
+    d = points[:, :, None, :] - boxes.centers.transpose(0, 1)[:, None, :, :]
+    c, s = torch.cos(boxes.yaw).T[:, None, :], torch.sin(boxes.yaw).T[:, None, :]
+    dx, dy, dz = d.unbind(-1)
+    return torch.stack((c * dx - s * dz, dy, s * dx + c * dz), dim=-1)
+
+
+def box_sdf(points: torch.Tensor, boxes: SceneBoxes) -> torch.Tensor:
+    """Signed distance (T, N, B) of points (T, N, 3) to each box; positive outside."""
+    q = to_box_local(points, boxes).abs() - boxes.half[None, None].to(device=points.device, dtype=points.dtype)
+    sq = torch.relu(q).square().sum(-1)
+    # sqrt has an infinite derivative at 0: route inside points through a constant.
+    outside = torch.where(sq > 0, torch.sqrt(torch.where(sq > 0, sq, torch.ones_like(sq))), torch.zeros_like(sq))
+    return outside + torch.clamp(q.max(-1).values, max=0)
+
+
+def floor_sdf(points: torch.Tensor, floor_y: float = 0.0) -> torch.Tensor:
+    return points[..., 1] - floor_y
+
+
+def penetration_loss(points: torch.Tensor, boxes: SceneBoxes, *, floor_y: float = 0.0, margin: float = 0.0) -> dict[str, torch.Tensor]:
+    """{"box", "floor"}: sum relu(margin - sd)^2 over points (T, N, 3) and every box / the floor."""
+    return {
+        "box": torch.relu(margin - box_sdf(points, boxes)).square().sum(),
+        "floor": torch.relu(margin - floor_sdf(points, floor_y)).square().sum(),
+    }
+
+
+@torch.no_grad()
+def penetration_stats(points: torch.Tensor, boxes: SceneBoxes, *, floor_y: float = 0.0) -> dict:
+    """Diagnostics: deepest box / floor penetration in cm and the frames with any penetration."""
+    box_depth = torch.relu(-box_sdf(points, boxes)).amax(dim=(1, 2)) if boxes.half.shape[0] else points.new_zeros(points.shape[0])
+    floor_depth = torch.relu(-floor_sdf(points, floor_y)).amax(dim=1)
+    worst = torch.maximum(box_depth, floor_depth)
+    return {
+        "maxBoxCm": float(box_depth.max()) * 100 if box_depth.numel() else 0.0,
+        "maxFloorCm": float(floor_depth.max()) * 100 if floor_depth.numel() else 0.0,
+        "frames": [int(t) for t in torch.nonzero(worst > 0).flatten()],
+    }
+
+
+@torch.no_grad()
+def support_height(points: torch.Tensor, boxes: SceneBoxes, *, floor_y: float = 0.0, step_tol: float = 0.1) -> torch.Tensor:
+    """(T, N) height of the surface under each point: the floor or the top of a box whose
+    footprint contains the point and whose top is at most `step_tol` above it."""
+    support = torch.full(points.shape[:2], float(floor_y), dtype=points.dtype, device=points.device)
+    if not boxes.half.shape[0]:
+        return support
+    boxes = _boxes_like(boxes, points)
+    local = to_box_local(points, boxes)
+    half = boxes.half[None, None]
+    over = (local[..., 0].abs() <= half[..., 0]) & (local[..., 2].abs() <= half[..., 2])
+    top = (boxes.centers[..., 1] + boxes.half[:, 1:2]).T[:, None, :]  # (T, 1, B)
+    eligible = over & (top <= points[..., 1:2] + step_tol)
+    tops = torch.where(eligible, top.expand_as(eligible), torch.full_like(eligible, -math.inf, dtype=points.dtype))
+    return torch.maximum(support, tops.amax(-1))
+
+
+def foot_points(joints: torch.Tensor) -> torch.Tensor:
+    """cskel27 joints (T, 27, 3) -> (T, 2, 2, 3): feet (left, right) x (ankle, toe)."""
+    return joints[:, torch.as_tensor(FOOT_JOINTS, device=joints.device)]
+
+
+@torch.no_grad()
+def stance_hmm(
+    feet,
+    fps: float,
+    boxes: SceneBoxes | None = None,
+    *,
+    floor_y: float = 0.0,
+    contact_height=None,
+    stay: float = 0.95,
+    height_scale: float = 0.05,
+    speed_scale: float = 0.3,
+    sharpness: float = 4.0,
+    step_tol: float = 0.1,
+) -> np.ndarray:
+    """Per-foot 2-state (swing/contact) Viterbi. Returns bool (T, F), True = contact.
+
+    feet: (T, F, 3) or (T, F, K, 3) world positions of the current estimate (K contact
+    joints per foot, e.g. `foot_points`). Per joint, h = height above its support minus
+    its planted height (`contact_height`, broadcastable to (F, K); default: the joint's
+    5th-percentile height over the clip), v = horizontal speed. Normalised evidence
+    z = min_k sqrt((max(h, 0) / height_scale)^2 + (v / speed_scale)^2) (any planted joint
+    means contact); p(contact) = sigmoid(sharpness * (1 - z)); stay probability `stay`.
+    """
+    # Discrete pass: everything (feet, boxes, contact_height) crosses to float64 CPU here.
+    x = _cpu64(feet)
+    if x.ndim == 3:
+        x = x[:, :, None, :]
+    if x.ndim != 4 or x.shape[-1] != 3:
+        raise SceneError("bad-feet", f"feet must be (T, F, 3) or (T, F, K, 3), got {tuple(x.shape)}")
+    if not torch.isfinite(x).all():
+        raise SceneError("non-finite", "foot positions contain NaN or inf")
+    t_len, f_len, k_len, _ = x.shape
+    flat = x.reshape(t_len, f_len * k_len, 3)
+    if boxes is None:
+        support = torch.full(flat.shape[:2], float(floor_y), dtype=torch.float64)
+    else:
+        support = support_height(flat, SceneBoxes(*(_cpu64(v) for v in boxes)), floor_y=floor_y, step_tol=step_tol)
+    height = (flat[..., 1] - support).reshape(t_len, f_len, k_len).numpy()
+    if contact_height is None:
+        ref = np.percentile(height, 5, axis=0)
+    else:
+        ref = np.broadcast_to(_cpu64(contact_height).numpy(), (f_len, k_len))
+    xz = x[..., [0, 2]].numpy()
+    speed = np.linalg.norm(np.gradient(xz, axis=0), axis=-1) * fps if t_len > 1 else np.zeros((t_len, f_len, k_len))
+    z = np.sqrt((np.maximum(height - ref, 0) / height_scale) ** 2 + (speed / speed_scale) ** 2).min(-1)  # (T, F)
+    logit = sharpness * (1 - z)
+    # -log sigmoid(l) = softplus(-l); cost[..., 0] swing, cost[..., 1] contact.
+    cost = np.stack((np.logaddexp(0, logit), np.logaddexp(0, -logit)), axis=-1)
+    switch = -math.log(1 - stay) + math.log(stay)
+    out = np.zeros((t_len, f_len), dtype=bool)
+    for f in range(f_len):
+        acc, back = cost[0, f].copy(), np.zeros((t_len, 2), dtype=np.int64)
+        for t in range(1, t_len):
+            nxt = np.empty(2)
+            for s in range(2):
+                via = acc + np.where(np.arange(2) == s, 0.0, switch)
+                back[t, s] = int(np.argmin(via))
+                nxt[s] = via[back[t, s]] + cost[t, f, s]
+            acc = nxt
+        state = int(np.argmin(acc))
+        for t in range(t_len - 1, -1, -1):
+            out[t, f] = state == 1
+            state = back[t, state]
+    return out
+
+
+def skate_loss(feet: torch.Tensor, stance, fps: float) -> torch.Tensor:
+    """Sum of squared horizontal velocity (m/s) of foot joints between consecutive
+    frames that are both stance. feet (T, F, 3) or (T, F, K, 3), stance bool (T, F)."""
+    if feet.ndim == 3:
+        feet = feet[:, :, None, :]
+    stance = torch.as_tensor(stance, dtype=torch.bool, device=feet.device)
+    both = (stance[1:] & stance[:-1])[..., None]  # (T-1, F, 1)
+    vel = (feet[1:, ..., [0, 2]] - feet[:-1, ..., [0, 2]]) * fps
+    return (vel.square().sum(-1) * both).sum()
+
+
+def camera_center(camera) -> np.ndarray:
+    """World camera position from camera.json (`worldToCamera`, OpenCV 4x4) or a 4x4 array."""
+    matrix = _cpu64(camera["worldToCamera"] if isinstance(camera, dict) else camera).numpy()
+    if matrix.shape != (4, 4) or not np.all(np.isfinite(matrix)):
+        raise SceneError("bad-camera", "worldToCamera must be a finite 4x4 matrix")
+    return -matrix[:3, :3].T @ matrix[:3, 3]
+
+
+@torch.no_grad()
+def ray_occlusion(points: torch.Tensor, camera_pos, boxes: SceneBoxes, *, eps: float = 1e-6) -> torch.Tensor:
+    """bool (T, N): the segment camera -> point passes through a box interior (the point is
+    hidden behind a box, or inside one). A point on a box's camera-facing face is visible."""
+    t_len, n = points.shape[:2]
+    if not boxes.half.shape[0]:
+        return torch.zeros((t_len, n), dtype=torch.bool, device=points.device)
+    boxes = _boxes_like(boxes, points)
+    cam = _cpu64(camera_pos).to(device=points.device, dtype=points.dtype)
+    origin = to_box_local(cam.expand(t_len, 1, 3), boxes)  # (T, 1, B, 3)
+    d = to_box_local(points, boxes) - origin  # (T, N, B, 3)
+    half = boxes.half[None, None].to(points.dtype)
+    flat = d.abs() < 1e-12
+    safe = torch.where(flat, torch.ones_like(d), d)
+    t0, t1 = (-half - origin) / safe, (half - origin) / safe
+    inside_slab = origin.abs() < half
+    lo = torch.where(flat, torch.where(inside_slab, -math.inf, math.inf), torch.minimum(t0, t1))
+    hi = torch.where(flat, torch.where(inside_slab, math.inf, -math.inf), torch.maximum(t0, t1))
+    near = torch.clamp(lo.amax(-1), min=0)
+    far = torch.clamp(hi.amin(-1), max=1)
+    return (far - near > eps).any(-1)

package/tools/track/py/test_joint_indices.py

@@ -0,0 +1,123 @@
+"""Independent name contracts for cskel27 versus COCO index spaces.
+
+The synthetic recovery fixture imports the objective's map, so it cannot by
+itself catch a shared off-by-one. These assertions use exported rig names and
+an independent anatomical correspondence instead of regenerating that map.
+"""
+import json
+from pathlib import Path
+import warnings
+
+import numpy as np
+import pytest
+import torch
+
+from lr_viterbi import ARM_PAIRS, COCO_NAMES, FACE_PAIRS, LEG_PAIRS, state_permutation
+from masks import ANKLES, EXTREMITIES, TORSO_PAIRS, UPPER
+from objective import (ARM_TRUST_JOINTS, COCO_JOINTS, JOINT_CAPS_DEG,
+                       articulation_reliability, trust_region_loss)
+from rig import State
+from scene import FOOT_JOINTS, FOOT_NAMES, foot_points
+
+
+@pytest.fixture(scope='module')
+def joint_names():
+    path = Path.home() / 'cclay-ingest/cozyfit/testdata/rig-y-bot-tpose.npz'
+    with np.load(path, allow_pickle=False) as rig:
+        meta = json.loads(bytes(rig['meta_json'].astype(np.uint8)))
+    names = meta['cskel27Joints']
+    assert len(names) == 27 and names[0] == 'Hips'
+    return names
+
+
+def test_coco_correspondence_and_visibility_indices_by_name(joint_names):
+    expected = ['LeftArm', 'RightArm', 'LeftForeArm', 'RightForeArm', 'LeftHand', 'RightHand',
+                'LeftUpLeg', 'RightUpLeg', 'LeftLeg', 'RightLeg', 'LeftFoot', 'RightFoot']
+    assert [joint_names[j] for j in COCO_JOINTS[5:]] == expected
+    # Visibility is anatomical: selecting it by this map must hide the matching
+    # detector-side body landmark, not the next joint along the right arm.
+    for coco, name in enumerate(expected, start=5):
+        hidden = torch.zeros(1, 27, dtype=torch.bool)
+        hidden[0, joint_names.index(name)] = True
+        assert hidden[:, COCO_JOINTS][0, 5:].nonzero().flatten().tolist() == [coco - 5]
+
+
+def test_trust_weighting_by_exported_bilateral_joint_names(joint_names):
+    strong = {side + part for side in ('Left', 'Right') for part in ('Shoulder', 'Arm', 'ForeArm', 'Hand')}
+    assert {joint_names[j] for j in ARM_TRUST_JOINTS} == strong
+    state = State.identity(1)
+    initial = state.local_rot_mats()
+    angle = 0.4
+    c, s = np.cos(angle), np.sin(angle)
+    turn = torch.tensor([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=torch.float32)
+    leg = initial.clone(); leg[:, joint_names.index('RightUpLeg')] = turn
+    unit = trust_region_loss(state.transl, leg, state.transl, initial).item()
+    assert unit > 0
+    for joint, name in enumerate(joint_names):
+        rotations = initial.clone(); rotations[:, joint] = turn
+        actual = trust_region_loss(state.transl, rotations, state.transl, initial).item()
+        multiplier = 2 if name == 'Hips' else 4 if name in strong else 1
+        assert actual == pytest.approx(multiplier * unit), name
+
+
+def test_joint_limit_caps_are_anatomically_symmetric(joint_names):
+    assert len(JOINT_CAPS_DEG) == len(joint_names)
+    caps = dict(zip(joint_names, JOINT_CAPS_DEG))
+    expected = {'Shoulder': 90, 'Arm': 175, 'ForeArm': 175, 'Hand': 110,
+                'HandEnd': 150, 'HandThumb1': 180, 'UpLeg': 160, 'Leg': 175, 'Foot': 100, 'ToeBase': 80}
+    for part, degrees in expected.items():
+        assert caps['Left' + part] == caps['Right' + part] == degrees
+
+
+def test_foot_and_lr_selections_use_their_own_index_spaces(joint_names):
+    assert FOOT_NAMES == ('left', 'right')
+    expected_feet = [['LeftFoot', 'LeftToeBase'], ['RightFoot', 'RightToeBase']]
+    assert [[joint_names[j] for j in foot] for foot in FOOT_JOINTS] == expected_feet
+    tagged = torch.arange(27)[None, :, None].expand(1, -1, 3)
+    selected = foot_points(tagged)[0, :, :, 0].tolist()
+    assert [[joint_names[j] for j in foot] for foot in selected] == expected_feet
+    for pairs, parts in [(ARM_PAIRS, ['Shoulder', 'Elbow', 'Wrist']),
+                         (LEG_PAIRS, ['Hip', 'Knee', 'Ankle']), (FACE_PAIRS, ['Eye', 'Ear'])]:
+        assert [(COCO_NAMES[a], COCO_NAMES[b]) for a, b in pairs] == [('left' + p, 'right' + p) for p in parts]
+    full = state_permutation('full_swap')
+    assert all(full[a] == b and full[b] == a for a, b in ARM_PAIRS + LEG_PAIRS + FACE_PAIRS)
+
+
+def test_mask_prompt_indices_by_coco_name():
+    assert [[COCO_NAMES[j] for j in pair] for pair in TORSO_PAIRS] == [
+        ['leftShoulder', 'rightShoulder'], ['leftHip', 'rightHip']]
+    assert [COCO_NAMES[j] for j in EXTREMITIES] == ['nose', 'leftWrist', 'rightWrist', 'leftAnkle', 'rightAnkle']
+    assert [COCO_NAMES[j] for j in UPPER] == ['nose', 'leftWrist', 'rightWrist']
+    assert [COCO_NAMES[j] for j in ANKLES] == ['leftAnkle', 'rightAnkle']
+
+
+def test_reliability_lower_median_strict_cuda_determinism():
+    assert torch.cuda.is_available(), 'Box acceptance requires CUDA'
+    projected = np.arange(4 * 17 * 2, dtype=np.float32).reshape(4, 17, 2) / 3
+    assigned = np.zeros((4, 17, 3), np.float32)
+    assigned[..., 2] = 1
+    assigned[:, :5, 2] = 0
+    assigned[1, 6::2, 2] = 0.2
+    assigned[2, 10:, 2] = 0
+    assigned[3, :, 2] = 0
+    seen = np.ones((4, 17), bool)
+    seen[2, 5] = False
+    expected = []
+    for uv, obs, visible in zip(projected, assigned, seen):
+        valid = visible[5:] & (obs[5:, 2] >= 0.5)
+        body = (uv - obs[:, :2])[5:][valid]
+        center = np.sort(body, axis=0)[(len(body) - 1) // 2] if len(body) else np.zeros(2)
+        excess = np.maximum(np.linalg.norm(uv - obs[:, :2] - center, axis=-1) / 3 - 1, 0) / 2
+        expected.append(obs[:, 2] * visible / (1 + excess ** 2) ** 2)
+    enabled = torch.are_deterministic_algorithms_enabled()
+    warn_only = torch.is_deterministic_algorithms_warn_only_enabled()
+    try:
+        torch.use_deterministic_algorithms(True)
+        with warnings.catch_warnings():
+            warnings.simplefilter('error', UserWarning)
+            actual = articulation_reliability(torch.tensor(projected, device='cuda'),
+                                              torch.tensor(assigned, device='cuda'),
+                                              torch.tensor(seen, device='cuda'), 3)
+        np.testing.assert_allclose(actual.cpu().numpy(), np.asarray(expected), atol=1e-6, rtol=1e-5)
+    finally:
+        torch.use_deterministic_algorithms(enabled, warn_only=warn_only)

package/tools/track/py/test_lr_viterbi.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+from eval_lr import evaluate_item, main as eval_main, truth_reference_error
+from lr_viterbi import (
+    ARM_PAIRS,
+    BODY_INDICES,
+    LEG_PAIRS,
+    STATE_NAMES,
+    solve_lr_viterbi,
+    state_permutation,
+)
+
+
+def synthetic_clip(frames: int = 64, seed: int = 17):
+    rng = np.random.default_rng(seed)
+    base = np.zeros((frames, 17, 2), dtype=np.float64)
+    # A moving torso plus distinct, smooth limbs.  The labels are deliberately
+    # geometric only; no colour or appearance signal is present.
+    for t in range(frames):
+        phase = 0.13 * t
+        base[t, 0] = [120 + 1.5 * t, 70]
+        base[t, 1:5] = [[119, 63], [121, 63], [117, 66], [123, 66]]
+        base[t, 5] = [100 + 2 * np.sin(phase), 120 + 0.3 * t]
+        base[t, 6] = [140 - 2 * np.sin(phase), 120 + 0.3 * t]
+        base[t, 7] = [82 + 3 * np.sin(phase), 155 + 2 * np.cos(phase)]
+        base[t, 8] = [158 - 3 * np.sin(phase), 155 + 2 * np.cos(phase)]
+        base[t, 9] = [70 + 4 * np.sin(phase), 190 + 2 * np.cos(phase)]
+        base[t, 10] = [170 - 4 * np.sin(phase), 190 + 2 * np.cos(phase)]
+        base[t, 11] = [108 + 1.5 * np.sin(phase), 220 + 0.4 * t]
+        base[t, 12] = [132 - 1.5 * np.sin(phase), 220 + 0.4 * t]
+        base[t, 13] = [102 + 2 * np.sin(phase), 270 + 1.5 * np.cos(phase)]
+        base[t, 14] = [138 - 2 * np.sin(phase), 270 + 1.5 * np.cos(phase)]
+        base[t, 15] = [96 + 2 * np.sin(phase), 320 + 1.2 * np.cos(phase)]
+        base[t, 16] = [144 - 2 * np.sin(phase), 320 + 1.2 * np.cos(phase)]
+    kp = np.concatenate([base, np.ones((frames, 17, 1), dtype=np.float64)], axis=2)
+    kp[..., :2] += rng.normal(0, 0.15, kp[..., :2].shape)
+    return base, kp
+
+
+def swap_observations(kp, groups, start, end):
+    out = kp.copy()
+    for t in range(start, end + 1):
+        for left, right in groups:
+            out[t, [left, right]] = out[t, [right, left]]
+    return out
+
+
+def assert_corrected(result, projected, observed, start, end, groups):
+    assignments = result["assignments"]
+    for t in range(projected.shape[0]):
+        for left, right in groups:
+            expected = right if start <= t <= end else left
+            assert assignments[t, left] == expected
+            expected = left if start <= t <= end else right
+            assert assignments[t, right] == expected
+    corrected = np.take_along_axis(observed[..., :2], assignments[:, :, None], axis=1)
+    raw_error = np.linalg.norm(observed[..., :2] - projected, axis=2)[:, list(BODY_INDICES)].mean()
+    corrected_error = np.linalg.norm(corrected - projected, axis=2)[:, list(BODY_INDICES)].mean()
+    # The detector noise remains after assignment correction; the swap error
+    # should nevertheless be removed rather than requiring a denoiser here.
+    assert corrected_error < raw_error * 0.4
+
+
+def test_one_three_seven_frame_swap_runs_are_recovered():
+    projected, clean = synthetic_clip()
+    # Use separate clips so each expected state is unambiguous.
+    for length, state_groups in ((1, ARM_PAIRS), (3, LEG_PAIRS), (7, ARM_PAIRS)):
+        start = 20
+        observed = swap_observations(clean, state_groups, start, start + length - 1)
+        result = solve_lr_viterbi(observed, projected, deltaPx=2.0)
+        assert_corrected(result, projected, observed, start, start + length - 1, state_groups)
+        expected_name = "arms_swap" if state_groups is ARM_PAIRS else "legs_swap"
+        assert all(result["state_names"][s] == expected_name for s in result["path"][start : start + length])
+
+
+def test_iid_noise_has_no_switches():
+    projected, clean = synthetic_clip(seed=21)
+    rng = np.random.default_rng(99)
+    noisy = clean.copy()
+    noisy[..., :2] += rng.normal(0, 1.0, noisy[..., :2].shape)
+    result = solve_lr_viterbi(noisy, projected, deltaPx=3.0)
+    assert result["switches"] == 0
+    assert np.all(result["path"] == STATE_NAMES.index("identity"))
+
+
+def test_slow_180_degree_turn_is_not_a_swap():
+    projected, _ = synthetic_clip(frames=80)
+    fixture = next((parent / "test/fixtures/heading-orient.json" for parent in Path(__file__).resolve().parents if (parent / "test/fixtures/heading-orient.json").exists()), None)
+    if fixture is not None:
+        raw = json.loads(fixture.read_text())["items"]["gt/turn"]
+        angles = []
+        for x, y, z in raw:
+            theta = np.linalg.norm([x, y, z])
+            qw = np.cos(theta / 2)
+            qy = (y / theta) * np.sin(theta / 2) if theta > 1e-12 else y / 2
+            angles.append(2 * np.arctan2(qy, qw))
+        angles = np.unwrap(np.asarray(angles))
+        projected = np.resize(projected, (len(angles), 17, 2))
+    else:
+        # The box runner receives only tools/track/py, so retain a deterministic
+        # equivalent when the repository fixture is not uploaded.
+        angles = np.linspace(0, np.pi, projected.shape[0])
+    turned = projected.copy()
+    for t, angle in enumerate(angles):
+        c, s = np.cos(angle), np.sin(angle)
+        xy = projected[t] - [125, 210]
+        turned[t] = xy @ np.array([[c, s], [-s, c]]) + [125, 210]
+    kp = np.concatenate([turned, np.ones((turned.shape[0], 17, 1))], axis=2)
+    kp[..., :2] += np.random.default_rng(123).normal(0, 0.2, kp[..., :2].shape)
+    result = solve_lr_viterbi(kp, turned, deltaPx=2.0)
+    assert result["switches"] == 0
+    assert np.all(result["path"] == STATE_NAMES.index("identity"))
+
+
+def test_malformed_and_no_evidence_inputs_are_safe():
+    projected = np.zeros((4, 17, 2), dtype=np.float64)
+    kp = np.zeros((4, 17, 3), dtype=np.float64)
+    result = solve_lr_viterbi(kp, projected, deltaPx=2.0)
+    assert np.all(result["path"] == 0)
+    assert np.all(result["ambiguous"])
+    assert np.all(result["margins"] == 0)
+
+    nan_kp = kp.copy()
+    nan_kp[..., :2] = np.nan
+    nan_result = solve_lr_viterbi(nan_kp, projected, deltaPx=2.0)
+    assert np.all(nan_result["path"] == 0)
+    empty = solve_lr_viterbi(np.empty((0, 17, 3)), np.empty((0, 17, 2)), deltaPx=2.0)
+    assert empty["path"].size == 0
+
+
+def test_permutations_are_restricted_and_deterministic():
+    assert np.array_equal(state_permutation("identity"), np.arange(17))
+    assert state_permutation("arms_swap")[11] == 11
+    assert state_permutation("legs_swap")[5] == 5
+    assert state_permutation("full_swap")[1] == 2
+    assert len(set(tuple(state_permutation(name)) for name in STATE_NAMES)) == 4
+
+
+def test_visibility_is_anatomical_not_permuted_with_detector_labels():
+    projected = np.zeros((1, 17, 2), dtype=np.float64)
+    observed = np.zeros((1, 17, 3), dtype=np.float64)
+    visibility = np.zeros((1, 17), dtype=bool)
+    # Only anatomical left shoulder is visible. Detector left is wrong by 100px;
+    # detector right is exactly at the visible anatomical projection.
+    projected[0, 5] = [0, 0]
+    observed[0, 5] = [100, 0, 1]
+    observed[0, 6] = [0, 0, 1]
+    visibility[0, 5] = True
+    paths = []
+    for hidden_right_x in (0, 10000):
+        projected[0, 6] = [hidden_right_x, 0]
+        result = solve_lr_viterbi(observed, projected, visibility=visibility, deltaPx=2.0)
+        paths.append(int(result["path"][0]))
+        assert int(np.argmin(result["emission"][0])) == STATE_NAMES.index("arms_swap")
+    assert paths == [STATE_NAMES.index("arms_swap")] * 2
+
+
+def test_hysteresis_keeps_clear_swap_when_unrelated_states_tie():
+    projected = np.zeros((5, 17, 2), dtype=np.float64)
+    observed = np.zeros((5, 17, 3), dtype=np.float64)
+    for left, right in ARM_PAIRS:
+        projected[:, right, 0] = 100
+        observed[:, left, 2] = observed[:, right, 2] = 1
+    observed[:, :, :2] = projected
+    for left, right in ARM_PAIRS:
+        observed[2, [left, right]] = observed[2, [right, left]]
+    result = solve_lr_viterbi(observed, projected, deltaPx=2.0, hysteresis=0.1)
+    assert result["raw_path"][2] in (STATE_NAMES.index("full_swap"), STATE_NAMES.index("arms_swap"))
+    assert result["assignments"][2, 5] == 6
+    assert result["assignments"][2, 6] == 5
+    assert result["margins"][2] > 0.1
+
+
+def test_eval_rejects_missing_truth_assignment_reference(tmp_path, monkeypatch):
+    item = {"set": "gt", "name": "walk", "variant": "shaded", "dir": str(tmp_path), "source": "none"}
+    assert truth_reference_error({}, "gt/walk") is not None
+    row = evaluate_item(item, {}, tmp_path, 2.0)
+    assert row["status"] == "missing-truth-reference"
+
+    summary = tmp_path / "summary.json"
+    approved = tmp_path / "approved.json"
+    out = tmp_path / "result.json"
+    summary.write_text(json.dumps({"deltaPx": 2.0}))
+    approved.write_text(json.dumps({"items": [item]}))
+    monkeypatch.setattr(sys, "argv", ["eval_lr.py", "--summary", str(summary), "--approved", str(approved), "--obs-root", str(tmp_path), "--out", str(out)])
+    assert eval_main() == 1
+    result = json.loads(out.read_text())
+    assert result["pass"] is False
+    assert result["items"][0]["status"] == "missing-truth-reference"