studio/rigging/overlay.py

"""studio.rigging.overlay — numpy + PIL debug overlay renderer.

Composites: source image → mask-edge ring → bones (cyan) → joints (red dots).
Output is a PIL.Image in RGBA, ready to save as PNG.
"""
from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw
from scipy.ndimage import binary_erosion

from pixel_cursor.rigging import BONES, JOINT_NAMES, Skeleton


JOINT_COLOR = (255, 80, 80, 255)
JOINT_OUTLINE = (0, 0, 0, 255)
BONE_COLOR = (60, 200, 255, 255)
MASK_EDGE_COLOR = (255, 255, 0, 200)
JOINT_RADIUS = 4
BONE_WIDTH = 2


def render_overlay(
    image: np.ndarray,
    mask: np.ndarray,
    skeleton: Skeleton,
    *,
    show_labels: bool = False,
) -> Image.Image:
    if image.ndim == 2:
        rgba = np.stack(
            [image, image, image, np.full_like(image, 255)],
            axis=-1,
        )
    elif image.shape[-1] == 3:
        rgba = np.concatenate(
            [image, np.full(image.shape[:2] + (1,), 255, dtype=np.uint8)],
            axis=-1,
        )
    elif image.shape[-1] == 4:
        rgba = image.copy()
    else:
        raise ValueError(f"unsupported image shape {image.shape}")
    H, W = rgba.shape[:2]

    mask_bool = mask.astype(bool)
    edge = mask_bool & ~binary_erosion(mask_bool, iterations=1, border_value=0)
    edge_layer_arr = np.zeros((H, W, 4), dtype=np.uint8)
    edge_layer_arr[edge] = MASK_EDGE_COLOR

    canvas = Image.alpha_composite(
        Image.fromarray(rgba, mode="RGBA"),
        Image.fromarray(edge_layer_arr, mode="RGBA"),
    )

    draw = ImageDraw.Draw(canvas)

    for parent_idx, child_idx in BONES:
        py, px = skeleton.positions[parent_idx]
        cy, cx = skeleton.positions[child_idx]
        draw.line(
            [(float(px), float(py)), (float(cx), float(cy))],
            fill=BONE_COLOR,
            width=BONE_WIDTH,
        )

    for i in range(skeleton.positions.shape[0]):
        y, x = skeleton.positions[i]
        r = JOINT_RADIUS
        draw.ellipse(
            [(float(x) - r, float(y) - r), (float(x) + r, float(y) + r)],
            fill=JOINT_COLOR,
            outline=JOINT_OUTLINE,
            width=1,
        )
        if show_labels:
            draw.text(
                (float(x) + r + 2, float(y) - r),
                JOINT_NAMES[i],
                fill=(255, 255, 255, 255),
            )

    return canvas


__all__ = ["render_overlay"]