Spaces:
Running
Running
File size: 5,370 Bytes
4e9daac | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """ControlNet-style visual previews for the Tagger tab.
Turns the YOLO pose keypoints produced by ``src.pose_tagger`` into previews a
video/animation pipeline can consume directly:
- ``skeleton_black`` — the pose skeleton on a black, letterboxed canvas
(drop-in OpenPose-style ControlNet input)
- ``skeleton_overlay`` — the skeleton drawn on top of the original image
- ``canny`` — Canny-style edge map computed in pure numpy (no
OpenCV / model downloads needed)
Keypoints are the COCO-17 arrays returned by ``PoseEstimator.estimate``
(already mapped back into original-image coordinates); rows are
``(x, y, confidence)``. Skeleton edges are confidence-gated at 0.25.
"""
from __future__ import annotations
import os
import tempfile
import numpy as np
from PIL import Image, ImageDraw
# COCO-17 skeleton edges (OpenPose ordering). Indices reference the keypoint
# order of ``src.pose_tagger``: nose, l/r eye, l/r ear, l/r shoulder, l/r
# elbow, l/r wrist, l/r hip, l/r knee, l/r ankle.
_SKELETON_EDGES = [
(0, 1), (0, 2), (1, 3), (2, 4), # face
(5, 6), (5, 7), (7, 9), (6, 8), (8, 10), # torso + arms
(5, 11), (6, 12), (11, 12), # hips
(11, 13), (13, 15), (12, 14), (14, 16), # legs
]
_KP_CONF = 0.25
_PALETTE = [
(255, 64, 128),
(0, 200, 255),
(140, 255, 64),
(255, 170, 0),
(170, 64, 255),
]
_KINDS = ("skeleton_black", "skeleton_overlay", "canny")
def _letterbox_scale(size, max_size: int) -> float:
ow, oh = size
if max_size <= 0:
return 1.0
return min(max_size / float(ow), max_size / float(oh)) if ow and oh else 1.0
def _visible(kp: np.ndarray) -> tuple[bool, ...]:
return tuple(bool(kp[i, 2] >= _KP_CONF) for i in range(len(kp)))
def _draw_skeleton(canvas: ImageDraw.ImageDraw, kpts: list[np.ndarray], scale: float = 1.0) -> None:
for pi, kp in enumerate(kpts):
vis = _visible(kp)
color = _PALETTE[pi % len(_PALETTE)]
for a, b in _SKELETON_EDGES:
if a >= len(kp) or b >= len(kp) or not (vis[a] and vis[b]):
continue
ax, ay = kp[a, 0] * scale, kp[a, 1] * scale
bx, by = kp[b, 0] * scale, kp[b, 1] * scale
canvas.line([ax, ay, bx, by], fill=color, width=3, joint="curve")
for i in range(len(kp)):
if not vis[i]:
continue
r = 4 if i <= 4 else 3
x, y = kp[i, 0] * scale, kp[i, 1] * scale
canvas.ellipse([x - r, y - r, x + r, y + r], fill=color)
def _skeleton_image(image: Image.Image, kpts: list[np.ndarray], black_bg: bool) -> Image.Image | None:
if not kpts:
return None
rgb = image.convert("RGB")
if black_bg:
scale = _letterbox_scale(rgb.size, 640)
cw, ch = max(1, round(rgb.width * scale)), max(1, round(rgb.height * scale))
canvas = Image.new("RGB", (cw, ch), (0, 0, 0))
else:
scale = 1.0
canvas = rgb.copy()
_draw_skeleton(ImageDraw.Draw(canvas), kpts, scale)
return canvas
def _convolve3(im: np.ndarray, kernel: np.ndarray) -> np.ndarray:
padded = np.pad(im, 1, mode="reflect")
out = np.zeros_like(im)
for i in range(3):
for j in range(3):
out += kernel[i, j] * padded[i:i + im.shape[0], j:j + im.shape[1]]
return out
def _canny_image(image: Image.Image, high: float = 0.35, low: float = 0.10) -> Image.Image:
gray = np.asarray(image.convert("L"), dtype=np.float64) / 255.0
blur = _convolve3(gray, np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=np.float64) / 16.0)
gx = _convolve3(blur, np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float64))
gy = _convolve3(blur, np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=np.float64))
mag = np.sqrt(gx * gx + gy * gy)
peak = float(mag.max())
if peak > 0:
mag /= peak
strong = mag >= high
weak = (mag >= low) & ~strong
dilated = strong.copy()
for dy in (-1, 0, 1):
for dx in (-1, 0, 1):
dilated |= np.roll(np.roll(strong, dy, axis=0), dx, axis=1)
edges = strong | (weak & dilated)
return Image.fromarray(np.where(edges, 255, 0).astype(np.uint8), mode="L").convert("RGB")
def render_control(image, keypoints: list[np.ndarray] | None, kind: str = "skeleton_black") -> Image.Image | None:
"""Render one ControlNet-style preview.
``keypoints`` is a list of COCO-17 arrays in image coordinates (as
returned by ``PoseEstimator.estimate``). Returns ``None`` when a skeleton
kind has no keypoints (caller hides the preview); the ``canny`` kind is
independent of pose.
"""
pil = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))
kpts = [np.asarray(k, dtype=np.float64) for k in (keypoints or [])]
kind = (kind or "skeleton_black").strip().lower()
if kind == "canny":
return _canny_image(pil)
if kind == "skeleton_overlay":
return _skeleton_image(pil, kpts, black_bg=False)
return _skeleton_image(pil, kpts, black_bg=True)
def save_control_png(pil: Image.Image) -> str:
"""Persist a preview to a temp file so ``gr.DownloadButton`` can serve it."""
fd, path = tempfile.mkstemp(prefix="whyx_control_", suffix=".png")
os.close(fd)
pil.save(path, "PNG")
return path
|