File size: 4,093 Bytes
7faaef2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """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",
]
|