| """Draw hand skeletons as 2D images. |
| |
| The same renderer is used for every embodiment and for humans, so the skeleton carries pose |
| only and never reveals what kind of hand produced it. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import cv2 |
| import numpy as np |
|
|
| from .topology import COLORS, EDGES, GROUP_OF |
|
|
|
|
| def draw_skeleton(kp: np.ndarray, vis: np.ndarray, width: int, height: int, |
| thickness: int = 2, radius: int = 3, supersample: int = 2) -> np.ndarray: |
| """Render one frame of hand skeletons on a black canvas. |
| |
| Args: |
| kp: ``(H, 17, 2)`` pixel coordinates. |
| vis: ``(H, 17)`` visibility. Invisible or non-finite keypoints are skipped together with |
| their edges. |
| width, height: canvas size in pixels. |
| |
| Returns: |
| ``(height, width, 3)`` uint8 RGB image. |
| """ |
| s = max(1, int(supersample)) |
| canvas = np.zeros((height * s, width * s, 3), np.uint8) |
| |
| |
| lw, dr = max(1, thickness * s // 2), max(1, radius * s // 2) |
| kp = np.asarray(kp, np.float64) |
| vis = np.asarray(vis, bool) & np.isfinite(kp).all(-1) |
| for h in range(kp.shape[0]): |
| pts = np.round(kp[h] * s).astype(np.int64) if vis[h].any() else None |
| if pts is None: |
| continue |
| for a, b, g in EDGES: |
| if vis[h, a] and vis[h, b]: |
| cv2.line(canvas, tuple(pts[a]), tuple(pts[b]), COLORS[g], lw, cv2.LINE_AA) |
| for j in range(kp.shape[1]): |
| if vis[h, j]: |
| cv2.circle(canvas, tuple(pts[j]), dr, COLORS[GROUP_OF[j]], -1, cv2.LINE_AA) |
| if s > 1: |
| canvas = cv2.resize(canvas, (width, height), interpolation=cv2.INTER_AREA) |
| return canvas |
|
|
|
|
| def draw_skeleton_video(kp: np.ndarray, vis: np.ndarray, width: int, height: int, |
| **kw) -> np.ndarray: |
| """``(T, H, 17, 2)`` keypoints -> ``(T, height, width, 3)`` uint8 skeleton video.""" |
| return np.stack([draw_skeleton(kp[t], vis[t], width, height, **kw) |
| for t in range(kp.shape[0])]) |
|
|
|
|
| def paint(frame: np.ndarray, skeleton: np.ndarray) -> np.ndarray: |
| """Overlay a rendered skeleton on a frame; skeleton pixels win wherever drawn.""" |
| out = frame.copy() |
| lit = skeleton.max(axis=-1) > 0 |
| out[lit] = skeleton[lit] |
| return out |
|
|
|
|
| def condition_frame(frame: np.ndarray, mask: np.ndarray | None, |
| skeleton: np.ndarray) -> np.ndarray: |
| """Phi(o; M, S) = ((1 - M) * o) (+) S: black out the mask region, then draw the skeleton. |
| |
| The skeleton is drawn even slightly outside the mask so fingertips are never clipped. |
| ``mask=None`` leaves the frame unmasked (mask-free inference). |
| """ |
| out = frame.copy() |
| if mask is not None: |
| out[np.asarray(mask).astype(bool)] = 0 |
| return paint(out, skeleton) |
|
|
