studio/rigging/joints.py

"""studio.rigging.joints — distance-transform + mask-topology joint detector.

Strategy: for each of 8 anatomical Y-fractions (head/neck/shoulder/elbow/wrist/
hip/knee/ankle), scan that row of the foreground mask, identify horizontal runs,
and map runs to joint indices. Distance-transform refines run centers to limb
midlines (avoids the "joint sits on the silhouette edge" failure mode).

Synthetic humanoids drawn with these same anatomical fractions verify the
detector within 8 px at every joint. Real cartoon sprites will need Session 2.
"""
from __future__ import annotations

from typing import List, Optional, Tuple

import numpy as np
from scipy.ndimage import distance_transform_edt

from pixel_cursor.rigging import NUM_JOINTS, Skeleton


ANATOMY_Y_FRACTIONS: dict[str, float] = {
    "head":     0.079,  # head center; mask top is head_top, not head_center
    "neck":     0.158,
    "shoulder": 0.189,
    "elbow":    0.333,
    "wrist":    0.482,
    "hip":      0.588,
    "knee":     0.789,
    "ankle":    0.991,
}


def _runs_in_row(row: np.ndarray) -> List[Tuple[int, int]]:
    if not row.any():
        return []
    padded = np.concatenate([[False], row, [False]])
    diff = np.diff(padded.astype(np.int8))
    starts = np.where(diff == 1)[0]
    ends = np.where(diff == -1)[0] - 1
    return list(zip(starts.tolist(), ends.tolist()))


def _refine_x_via_dt(dt_row: np.ndarray, x_start: int, x_end: int) -> float:
    """Within a horizontal run, return the X with maximum distance-transform value."""
    segment = dt_row[x_start : x_end + 1]
    best = int(np.argmax(segment))
    return float(x_start + best)


def _pick_lr_runs(
    runs: List[Tuple[int, int]],
) -> Optional[Tuple[Tuple[int, int], Tuple[int, int]]]:
    if len(runs) < 2:
        return None
    if len(runs) == 2:
        a, b = sorted(runs, key=lambda r: r[0])
        return a, b
    sorted_runs = sorted(runs, key=lambda r: r[0])
    return sorted_runs[0], sorted_runs[-1]


def detect_joints(mask: np.ndarray) -> Skeleton:
    if mask.ndim != 2:
        raise ValueError(f"mask must be 2D, got shape {mask.shape}")
    mask_bool = mask.astype(bool)
    H, W = mask_bool.shape

    ys, _ = np.where(mask_bool)
    if ys.size == 0:
        return Skeleton(
            positions=np.zeros((NUM_JOINTS, 2), dtype=np.float32),
            confidence=np.zeros(NUM_JOINTS, dtype=np.float32),
            image_shape=(H, W),
        )

    y_min, y_max = int(ys.min()), int(ys.max())
    bbox_h = y_max - y_min + 1
    dt = distance_transform_edt(mask_bool).astype(np.float32)

    positions = np.zeros((NUM_JOINTS, 2), dtype=np.float32)
    confidence = np.ones(NUM_JOINTS, dtype=np.float32)

    def anatomy_y(name: str) -> int:
        return int(np.clip(y_min + ANATOMY_Y_FRACTIONS[name] * bbox_h, y_min, y_max))

    def single_x(y: int) -> Tuple[float, float]:
        runs = _runs_in_row(mask_bool[y])
        if not runs:
            return float(W) / 2.0, 0.0
        s, e = max(runs, key=lambda r: r[1] - r[0])
        return _refine_x_via_dt(dt[y], s, e), 1.0

    def lr_x(y: int) -> Optional[Tuple[float, float]]:
        runs = _runs_in_row(mask_bool[y])
        pair = _pick_lr_runs(runs)
        if pair is None:
            return None
        (ls, le), (rs, re) = pair
        return _refine_x_via_dt(dt[y], ls, le), _refine_x_via_dt(dt[y], rs, re)

    head_y = anatomy_y("head")
    hx, hconf = single_x(head_y)
    positions[0] = (head_y, hx)
    confidence[0] = hconf

    neck_y = anatomy_y("neck")
    nx, nconf = single_x(neck_y)
    positions[1] = (neck_y, nx)
    confidence[1] = nconf

    shoulder_y = anatomy_y("shoulder")
    lr = lr_x(shoulder_y)
    if lr is not None:
        positions[2] = (shoulder_y, lr[0])
        positions[3] = (shoulder_y, lr[1])
    else:
        runs = _runs_in_row(mask_bool[shoulder_y])
        if runs:
            s, e = max(runs, key=lambda r: r[1] - r[0])
            positions[2] = (shoulder_y, s)
            positions[3] = (shoulder_y, e)
            confidence[2] = confidence[3] = 0.5
        else:
            confidence[2] = confidence[3] = 0.0

    for joint_l, joint_r, key in ((4, 5, "elbow"), (6, 7, "wrist")):
        y = anatomy_y(key)
        lr = lr_x(y)
        if lr is not None:
            positions[joint_l] = (y, lr[0])
            positions[joint_r] = (y, lr[1])
        else:
            positions[joint_l] = (y, positions[2, 1])
            positions[joint_r] = (y, positions[3, 1])
            confidence[joint_l] = confidence[joint_r] = 0.5

    hip_y = anatomy_y("hip")
    lr = lr_x(hip_y)
    if lr is not None:
        positions[8] = (hip_y, lr[0])
        positions[9] = (hip_y, lr[1])
    else:
        runs = _runs_in_row(mask_bool[hip_y])
        if runs:
            s, e = max(runs, key=lambda r: r[1] - r[0])
            positions[8] = (hip_y, s)
            positions[9] = (hip_y, e)
            confidence[8] = confidence[9] = 0.5
        else:
            confidence[8] = confidence[9] = 0.0

    for joint_l, joint_r, key in ((10, 11, "knee"), (12, 13, "ankle")):
        y = anatomy_y(key)
        lr = lr_x(y)
        if lr is not None:
            positions[joint_l] = (y, lr[0])
            positions[joint_r] = (y, lr[1])
        else:
            positions[joint_l] = (y, positions[8, 1])
            positions[joint_r] = (y, positions[9, 1])
            confidence[joint_l] = confidence[joint_r] = 0.5

    return Skeleton(positions=positions, confidence=confidence, image_shape=(H, W))


__all__ = ["detect_joints", "ANATOMY_Y_FRACTIONS"]