Spaces:
Running
Running
| """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 | |