studio/rigging/skin.py

"""studio.rigging.skin — inverse-distance-falloff mesh skinning weights.

For each foreground pixel and each bone, weight ∝ 1 / distance_to_bone^p,
normalised across bones to sum to 1 per pixel. Background pixels get zero
weight on every bone.

This is the "fast / good-enough" baseline picked in PHASE_3_RESEARCH.md.
Heat-diffusion skinning (better crease quality) is the upgrade path for
Session 4 if creases hurt visual quality.
"""
from __future__ import annotations

import numpy as np

from pixel_cursor.rigging import BONES, NUM_JOINTS, Skeleton


DEFAULT_FALLOFF: float = 4.0
DEFAULT_EPSILON: float = 0.5


def _point_segment_distance(
    points: np.ndarray, a: np.ndarray, b: np.ndarray
) -> np.ndarray:
    """Distance from each point to the closed segment AB.

    points: shape (..., 2) in (y, x)
    a, b:   shape (2,) endpoints
    Returns shape (...)
    """
    ab = b - a
    ap = points - a
    ab_len_sq = float((ab * ab).sum())
    if ab_len_sq < 1e-6:
        return np.linalg.norm(ap, axis=-1)
    t = (ap * ab).sum(axis=-1) / ab_len_sq
    t = np.clip(t, 0.0, 1.0)
    closest = a + t[..., None] * ab
    return np.linalg.norm(points - closest, axis=-1)


def compute_skin_weights(
    skeleton: Skeleton,
    mask: np.ndarray,
    *,
    falloff: float = DEFAULT_FALLOFF,
    epsilon: float = DEFAULT_EPSILON,
) -> np.ndarray:
    """Return (H, W, N_BONES) inverse-distance-falloff skinning weights.

    Properties:
      - Each foreground pixel's weights sum to 1.0 (partition of unity).
      - Each background pixel's weights are exactly 0 across all bones.
      - Deterministic: same skeleton + mask → same weights, bitwise.
    """
    if mask.ndim != 2:
        raise ValueError(f"mask must be 2D, got {mask.shape}")
    if skeleton.image_shape != mask.shape:
        raise ValueError(
            f"skeleton.image_shape={skeleton.image_shape} != mask.shape={mask.shape}"
        )

    H, W = mask.shape
    ys, xs = np.indices((H, W))
    coords = np.stack([ys, xs], axis=-1).astype(np.float32)

    n_bones = len(BONES)
    dists = np.empty((H, W, n_bones), dtype=np.float32)
    rest = skeleton.positions
    for b, (parent_idx, child_idx) in enumerate(BONES):
        dists[..., b] = _point_segment_distance(
            coords, rest[parent_idx], rest[child_idx]
        )

    inv = 1.0 / np.maximum(dists, epsilon) ** falloff
    total = inv.sum(axis=-1, keepdims=True)
    weights = inv / np.maximum(total, 1e-12)

    mask_bool = mask.astype(bool)[..., None]
    weights = np.where(mask_bool, weights, 0.0).astype(np.float32)
    return weights


__all__ = ["compute_skin_weights", "DEFAULT_FALLOFF", "DEFAULT_EPSILON"]