pixel_cursor/rigging.py

"""pixel_cursor.rigging — 14-joint humanoid skeleton primitives.

Session 1 (Phase 3) scope: data type, canonical joint layout, parent-linkage.
Detection / segmentation / overlay live in studio.rigging.

The skeleton is image-space (pixel coordinates, (y, x) order). Coordinates are
floats so sub-pixel joint positions (from distance-transform peak interpolation)
survive without rounding loss.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Final, Tuple

import numpy as np


JOINT_NAMES: Final[Tuple[str, ...]] = (
    "head",        # 0
    "neck",        # 1  — root (parent = -1)
    "l_shoulder",  # 2
    "r_shoulder",  # 3
    "l_elbow",     # 4
    "r_elbow",     # 5
    "l_wrist",     # 6
    "r_wrist",     # 7
    "l_hip",       # 8
    "r_hip",       # 9
    "l_knee",      # 10
    "r_knee",      # 11
    "l_ankle",     # 12
    "r_ankle",     # 13
)

NUM_JOINTS: Final[int] = len(JOINT_NAMES)

PARENTS: Final[Tuple[int, ...]] = (
    1,   # head       -> neck
    -1,  # neck       -> root
    1,   # l_shoulder -> neck
    1,   # r_shoulder -> neck
    2,   # l_elbow    -> l_shoulder
    3,   # r_elbow    -> r_shoulder
    4,   # l_wrist    -> l_elbow
    5,   # r_wrist    -> r_elbow
    1,   # l_hip      -> neck
    1,   # r_hip      -> neck
    8,   # l_knee     -> l_hip
    9,   # r_knee     -> r_hip
    10,  # l_ankle    -> l_knee
    11,  # r_ankle    -> r_knee
)

JOINT_INDEX: Final[dict[str, int]] = {name: i for i, name in enumerate(JOINT_NAMES)}

BONES: Final[Tuple[Tuple[int, int], ...]] = tuple(
    (PARENTS[i], i) for i in range(NUM_JOINTS) if PARENTS[i] >= 0
)


@dataclass(frozen=True, eq=False, repr=False)
class Skeleton:
    positions: np.ndarray         # shape (14, 2), float32, (y, x) per joint
    confidence: np.ndarray         # shape (14,),  float32, in [0, 1]
    image_shape: Tuple[int, int]   # (H, W) of the source mask

    def __post_init__(self) -> None:
        if self.positions.shape != (NUM_JOINTS, 2):
            raise ValueError(
                f"positions must be shape ({NUM_JOINTS}, 2), got {self.positions.shape}"
            )
        if self.confidence.shape != (NUM_JOINTS,):
            raise ValueError(
                f"confidence must be shape ({NUM_JOINTS},), got {self.confidence.shape}"
            )
        if self.positions.dtype != np.float32:
            object.__setattr__(self, "positions", self.positions.astype(np.float32))
        if self.confidence.dtype != np.float32:
            object.__setattr__(self, "confidence", self.confidence.astype(np.float32))
        # Lock buffers so the frozen contract survives array views.
        self.positions.setflags(write=False)
        self.confidence.setflags(write=False)

    def __repr__(self) -> str:
        return (
            f"Skeleton(image_shape={self.image_shape}, "
            f"joints={NUM_JOINTS}, "
            f"mean_confidence={float(self.confidence.mean()):.2f})"
        )

    @classmethod
    def from_positions(
        cls,
        positions: np.ndarray,
        image_shape: Tuple[int, int],
        confidence: np.ndarray | None = None,
    ) -> "Skeleton":
        positions = np.asarray(positions, dtype=np.float32)
        if confidence is None:
            confidence = np.ones(NUM_JOINTS, dtype=np.float32)
        else:
            confidence = np.asarray(confidence, dtype=np.float32)
        return cls(positions=positions, confidence=confidence, image_shape=image_shape)

    def joint(self, name: str) -> Tuple[float, float]:
        idx = JOINT_INDEX[name]
        return float(self.positions[idx, 0]), float(self.positions[idx, 1])

    def to_dict(self) -> dict[str, dict[str, float]]:
        return {
            name: {
                "y": float(self.positions[i, 0]),
                "x": float(self.positions[i, 1]),
                "confidence": float(self.confidence[i]),
            }
            for i, name in enumerate(JOINT_NAMES)
        }


__all__ = [
    "JOINT_NAMES",
    "NUM_JOINTS",
    "PARENTS",
    "JOINT_INDEX",
    "BONES",
    "Skeleton",
]