File size: 2,403 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 | """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"]
|