utils/skeleton.py

"""
utils/skeleton.py
─────────────────
Bone registry, parent/child hierarchy, H36M↔custom-name mapping,
T-pose rest positions, and all coordinate-space helpers.

Custom bone layout
──────────────────
Idx  Name             Parent  H36M src joint
  0  Hips                -1    0  (hip centre)
  1  Spine                0    7  (spine)
  2  Chest                1    8  (thorax)
  3  Neck                 2    lerp(chest, head, 0.4)
  4  Head                 3   10  (head / top)
  5  LeftUpperArm         2   11  (L shoulder)
  6  LeftLowerArm         5   12  (L elbow)
  7  LeftHand             6   13  (L wrist)
  8  RightUpperArm        2   14  (R shoulder)
  9  RightLowerArm        8   15  (R elbow)
 10  RightHand            9   16  (R wrist)
 11  LeftUpperLeg         0    4  (L hip)
 12  LeftLowerLeg        11    5  (L knee)
 13  LeftFoot            12    6  (L ankle)
 14  RightUpperLeg        0    1  (R hip)
 15  RightLowerLeg       14    2  (R knee)
 16  RightFoot           15    3  (R ankle)
"""

from __future__ import annotations
import numpy as np

# ── Registry ──────────────────────────────────────────────────────────────────

BONE_NAMES: list[str] = [
    "Hips",           #  0
    "Spine",          #  1
    "Chest",          #  2
    "Neck",           #  3
    "Head",           #  4
    "LeftUpperArm",   #  5
    "LeftLowerArm",   #  6
    "LeftHand",       #  7
    "RightUpperArm",  #  8
    "RightLowerArm",  #  9
    "RightHand",      # 10
    "LeftUpperLeg",   # 11
    "LeftLowerLeg",   # 12
    "LeftFoot",       # 13
    "RightUpperLeg",  # 14
    "RightLowerLeg",  # 15
    "RightFoot",      # 16
]

N_BONES: int = len(BONE_NAMES)
BONE_IDX: dict[str, int] = {n: i for i, n in enumerate(BONE_NAMES)}

BONE_PARENT: list[int] = [
    -1,  #  0  Hips         root
     0,  #  1  Spine        ← Hips
     1,  #  2  Chest        ← Spine
     2,  #  3  Neck         ← Chest
     3,  #  4  Head         ← Neck
     2,  #  5  LeftUpperArm ← Chest
     5,  #  6  LeftLowerArm ← LeftUpperArm
     6,  #  7  LeftHand     ← LeftLowerArm
     2,  #  8  RightUpperArm← Chest
     8,  #  9  RightLowerArm← RightUpperArm
     9,  # 10  RightHand    ← RightLowerArm
     0,  # 11  LeftUpperLeg ← Hips
    11,  # 12  LeftLowerLeg ← LeftUpperLeg
    12,  # 13  LeftFoot     ← LeftLowerLeg
     0,  # 14  RightUpperLeg← Hips
    14,  # 15  RightLowerLeg← RightUpperLeg
    15,  # 16  RightFoot    ← RightLowerLeg
]

# Derived child lists
BONE_CHILDREN: dict[int, list[int]] = {i: [] for i in range(N_BONES)}
for _c, _p in enumerate(BONE_PARENT):
    if _p >= 0:
        BONE_CHILDREN[_p].append(_c)

# ── T-Pose rest positions (world space, metres, Y-up) ─────────────────────────
# Used for GLB inverse-bind matrices and Godot rest pose.
T_POSE_WORLD: np.ndarray = np.array([
    [ 0.000,  1.000,  0.000],  #  0  Hips
    [ 0.000,  1.100,  0.000],  #  1  Spine
    [ 0.000,  1.300,  0.000],  #  2  Chest
    [ 0.000,  1.450,  0.000],  #  3  Neck
    [ 0.000,  1.620,  0.000],  #  4  Head
    [-0.200,  1.350,  0.000],  #  5  LeftUpperArm
    [-0.450,  1.350,  0.000],  #  6  LeftLowerArm
    [-0.650,  1.350,  0.000],  #  7  LeftHand
    [ 0.200,  1.350,  0.000],  #  8  RightUpperArm
    [ 0.450,  1.350,  0.000],  #  9  RightLowerArm
    [ 0.650,  1.350,  0.000],  # 10  RightHand
    [-0.100,  0.850,  0.000],  # 11  LeftUpperLeg
    [-0.100,  0.450,  0.000],  # 12  LeftLowerLeg
    [-0.100,  0.050,  0.000],  # 13  LeftFoot
    [ 0.100,  0.850,  0.000],  # 14  RightUpperLeg
    [ 0.100,  0.450,  0.000],  # 15  RightLowerLeg
    [ 0.100,  0.050,  0.000],  # 16  RightFoot
], dtype=np.float32)

# T-pose LOCAL positions (relative to parent bone)
def _build_tpose_local() -> np.ndarray:
    local = np.zeros_like(T_POSE_WORLD)
    for j in range(N_BONES):
        p = BONE_PARENT[j]
        if p < 0:
            local[j] = T_POSE_WORLD[j]
        else:
            local[j] = T_POSE_WORLD[j] - T_POSE_WORLD[p]
    return local

T_POSE_LOCAL: np.ndarray = _build_tpose_local()

# ── H36M joint-index → our bone index ────────────────────────────────────────
# H36M 17-joint order:
#  0=hip_ctr, 1=R_hip, 2=R_knee, 3=R_ankle,
#  4=L_hip,  5=L_knee, 6=L_ankle,
#  7=spine,  8=thorax, 9=nose,  10=head,
# 11=L_shoulder,12=L_elbow,13=L_wrist,
# 14=R_shoulder,15=R_elbow,16=R_wrist

H36M_TO_BONE: dict[int, int] = {
     0:  0,   # hip ctr    → Hips
     7:  1,   # spine      → Spine
     8:  2,   # thorax     → Chest
     # 9 (nose) → Neck  [computed as lerp]
    10:  4,   # head       → Head
    11:  5,   # L shoulder → LeftUpperArm
    12:  6,   # L elbow    → LeftLowerArm
    13:  7,   # L wrist    → LeftHand
    14:  8,   # R shoulder → RightUpperArm
    15:  9,   # R elbow    → RightLowerArm
    16: 10,   # R wrist    → RightHand
     4: 11,   # L hip      → LeftUpperLeg
     5: 12,   # L knee     → LeftLowerLeg
     6: 13,   # L ankle    → LeftFoot
     1: 14,   # R hip      → RightUpperLeg
     2: 15,   # R knee     → RightLowerLeg
     3: 16,   # R ankle    → RightFoot
}

# ── COCO-17 → H36M-17 (2-D) ──────────────────────────────────────────────────
# COCO: 0=nose,1=L_eye,2=R_eye,3=L_ear,4=R_ear,5=L_sh,6=R_sh,
#       7=L_el,8=R_el,9=L_wr,10=R_wr,11=L_hip,12=R_hip,
#       13=L_kn,14=R_kn,15=L_an,16=R_an

def coco_to_h36m_2d(coco: np.ndarray) -> np.ndarray:
    """
    (T, 17, 2) COCO pixel coords  →  (T, 17, 2) H36M-layout pixel coords.
    """
    T = coco.shape[0]
    h = np.zeros((T, 17, 2), dtype=np.float32)

    # Direct copies
    h[:, 1,  :] = coco[:, 12, :]   # R_hip
    h[:, 2,  :] = coco[:, 14, :]   # R_knee
    h[:, 3,  :] = coco[:, 16, :]   # R_ankle
    h[:, 4,  :] = coco[:, 11, :]   # L_hip
    h[:, 5,  :] = coco[:, 13, :]   # L_knee
    h[:, 6,  :] = coco[:, 15, :]   # L_ankle
    h[:, 9,  :] = coco[:, 0,  :]   # nose  → H36M nose slot
    h[:, 11, :] = coco[:, 5,  :]   # L_shoulder
    h[:, 12, :] = coco[:, 7,  :]   # L_elbow
    h[:, 13, :] = coco[:, 9,  :]   # L_wrist
    h[:, 14, :] = coco[:, 6,  :]   # R_shoulder
    h[:, 15, :] = coco[:, 8,  :]   # R_elbow
    h[:, 16, :] = coco[:, 10, :]   # R_wrist

    # Derived midpoints
    h[:, 0,  :] = (coco[:, 11, :] + coco[:, 12, :]) * 0.5  # hip centre
    h[:, 8,  :] = (coco[:, 5,  :] + coco[:, 6,  :]) * 0.5  # thorax (shoulder mid)
    h[:, 7,  :] = (h[:, 0, :] + h[:, 8, :]) * 0.5           # spine midpoint
    h[:, 10, :] = (coco[:, 3,  :] + coco[:, 4,  :]) * 0.5  # head (ear mid)

    return h


def h36m_to_bone_positions(h36m_world: np.ndarray) -> np.ndarray:
    """
    Map H36M (T, 17, 3) world positions → our bone layout (T, N_BONES, 3).

    Joint 3 (Neck) is linearly interpolated between Chest (bone 2) and Head (bone 4).
    All other bones have a direct H36M source.
    """
    T = h36m_world.shape[0]
    out = np.zeros((T, N_BONES, 3), dtype=np.float32)

    for h36m_j, bone_j in H36M_TO_BONE.items():
        out[:, bone_j, :] = h36m_world[:, h36m_j, :]

    # Neck = 60 % Chest + 40 % Head
    out[:, 3, :] = out[:, 2, :] * 0.60 + out[:, 4, :] * 0.40

    return out


def compute_local_positions(world: np.ndarray) -> np.ndarray:
    """
    (T, N_BONES, 3) world positions → (T, N_BONES, 3) parent-relative positions.
    Root (Hips) keeps its world position.
    """
    local = np.empty_like(world)
    local[:, 0, :] = world[:, 0, :]
    for j in range(1, N_BONES):
        local[:, j, :] = world[:, j, :] - world[:, BONE_PARENT[j], :]
    return local