| """Styled pose-skeleton renderer for AnimaPoseControl. Pure numpy + cv2 (no torch). |
| Single source of truth for the predefined skeleton looks. |
| Wholebody-133: body 0-16, feet 17-22, face 23-90, L-hand 91-111, R-hand 112-132.""" |
| import numpy as np |
| import cv2 |
|
|
| BODY = list(range(17)); FEET = list(range(17, 23)); FACE = list(range(23, 91)); HAND = list(range(91, 133)) |
| COCO_LIMBS = [(5, 7), (7, 9), (6, 8), (8, 10), (11, 13), (13, 15), (12, 14), (14, 16), |
| (5, 6), (11, 12), (5, 11), (6, 12), (0, 5), (0, 6)] |
| _PAL = [(255, 0, 0), (255, 128, 0), (255, 255, 0), (128, 255, 0), (0, 255, 0), (0, 255, 128), (0, 255, 255), |
| (0, 128, 255), (0, 0, 255), (128, 0, 255), (255, 0, 255), (255, 0, 128), (180, 180, 180), (220, 220, 220)] |
| LIMB_COLOR = {l: _PAL[i % len(_PAL)] for i, l in enumerate(COCO_LIMBS)} |
| JOINT_COLOR = [_PAL[i % len(_PAL)] for i in range(17)] |
| STYLES = ("R0_thin", "R1_thick", "R2_puppet", "heatmap") |
| _W = {"R0_thin": {"stick": 2, "joint": 3}, "R1_thick": {"stick": 8, "joint": 6}, "R2_puppet": {"stick": 22, "joint": 11}, "heatmap": {"stick": 2, "joint": 3}} |
|
|
|
|
| def _toggle(sc, hands, face, feet): |
| sc = np.asarray(sc, float).copy() |
| if not hands: sc[HAND] = 0.0 |
| if not face: sc[FACE] = 0.0 |
| if not feet: sc[FEET] = 0.0 |
| return sc |
|
|
|
|
| def render(kp, sc, style, resolution, hands=True, face=True, feet=True, kpt_thr=0.3): |
| if style not in STYLES: |
| raise ValueError(f"unknown style {style!r}") |
| H = W = int(resolution) |
| kp = np.asarray(kp, float); sc = _toggle(sc, hands, face, feet) |
| if style == "heatmap": |
| return _heatmap(kp, sc, H, W, feet, kpt_thr) |
| cfg = _W[style]; cv = np.zeros((H, W, 3), np.uint8) |
| for a, b in COCO_LIMBS: |
| if sc[a] >= kpt_thr and sc[b] >= kpt_thr: |
| cv2.line(cv, (int(kp[a][0]), int(kp[a][1])), (int(kp[b][0]), int(kp[b][1])), LIMB_COLOR[(a, b)], cfg["stick"]) |
| for i in BODY: |
| if sc[i] >= kpt_thr: |
| cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), cfg["joint"], (255, 255, 255), -1) |
| for i in FEET: |
| if sc[i] >= kpt_thr: |
| cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), max(1, cfg["joint"] // 2), (255, 255, 255), -1) |
| if style == "R0_thin": |
| for i in HAND: |
| if sc[i] >= kpt_thr: |
| cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), 2, (0, 255, 255), -1) |
| for i in FACE: |
| if sc[i] >= kpt_thr: |
| cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), 2, (255, 255, 255), -1) |
| return cv |
|
|
|
|
| def _heatmap(kp, sc, H, W, feet, thr, sig=11): |
| acc = np.zeros((H, W, 3), np.float32) |
| for i in BODY + (FEET if feet else []): |
| if sc[i] >= thr: |
| x, y = int(round(kp[i][0])), int(round(kp[i][1])) |
| if 0 <= x < W and 0 <= y < H: |
| b = np.zeros((H, W), np.float32); b[y, x] = 1.0 |
| b = cv2.GaussianBlur(b, (0, 0), sig); b /= (b.max() + 1e-9) |
| acc += b[..., None] * np.array(JOINT_COLOR[i % 17], np.float32) |
| return np.clip(acc, 0, 255).astype(np.uint8) |
|
|
