458 lines
16 KiB
Python
458 lines
16 KiB
Python
"""
|
|
OpenPose keypoint parsing, torso detection, 3D rotation, and skeleton rendering.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
# Body keypoint indices (0-based, OpenPose COCO format)
|
|
# 0=nose, 1=neck, 2=r_shoulder, 3=r_elbow, 4=r_wrist, 5=l_shoulder, 6=l_elbow, 7=l_wrist,
|
|
# 8=mid_hip, 9=r_hip, 10=r_knee, 11=r_ankle, 12=l_hip, 13=l_knee, 14=l_ankle,
|
|
# 15=r_eye, 16=l_eye, 17=r_ear, 18=l_ear
|
|
TORSO_INDICES = [1, 2, 5, 8, 9, 12] # neck, shoulders, mid_hip, hips
|
|
|
|
# Face/head keypoints (hidden when rotated >90° - back of head visible)
|
|
FACE_INDICES = {0, 15, 16, 17, 18}
|
|
|
|
# Limb-specific depth scales (OpenPose COCO body indices 0-18)
|
|
# Torso compact; arms/legs extend further; head moderate
|
|
# Based on typical human proportions in frontal view
|
|
BODY_DEPTH_SCALES: dict[int, float] = {
|
|
0: 0.30, # nose - head
|
|
1: 0.20, # neck - torso center
|
|
2: 0.35, # r_shoulder
|
|
3: 0.50, # r_elbow - arm extends
|
|
4: 0.55, # r_wrist
|
|
5: 0.35, # l_shoulder
|
|
6: 0.50, # l_elbow
|
|
7: 0.55, # l_wrist
|
|
8: 0.22, # mid_hip - torso
|
|
9: 0.30, # r_hip
|
|
10: 0.45, # r_knee - leg extends
|
|
11: 0.50, # r_ankle
|
|
12: 0.30, # l_hip
|
|
13: 0.45, # l_knee
|
|
14: 0.50, # l_ankle
|
|
15: 0.28, # r_eye
|
|
16: 0.28, # l_eye
|
|
17: 0.30, # r_ear
|
|
18: 0.30, # l_ear
|
|
}
|
|
|
|
# Default depth scale for hands (21 keypoints each - no per-index in OpenPose hand spec)
|
|
HAND_DEPTH_SCALE = 0.45
|
|
|
|
# Occlusion: points with depth below this (behind torso) get confidence zeroed
|
|
OCCLUSION_DEPTH_THRESHOLD = -0.15
|
|
|
|
# Limb connections for body (1-indexed in ControlNet util, we use 0-indexed)
|
|
# limbSeq: (from_part, to_part) - 0-indexed
|
|
BODY_LIMBS = [
|
|
(1, 2), # neck - r_shoulder
|
|
(1, 5), # neck - l_shoulder
|
|
(2, 3), # r_shoulder - r_elbow
|
|
(3, 4), # r_elbow - r_wrist
|
|
(5, 6), # l_shoulder - l_elbow
|
|
(6, 7), # l_elbow - l_wrist
|
|
(1, 8), # neck - mid_hip
|
|
(8, 9), # mid_hip - r_hip
|
|
(9, 10), # r_hip - r_knee
|
|
(10, 11), # r_knee - r_ankle
|
|
(8, 12), # mid_hip - l_hip
|
|
(12, 13), # l_hip - l_knee
|
|
(13, 14), # l_knee - l_ankle
|
|
(0, 1), # nose - neck
|
|
(0, 15), # nose - r_eye
|
|
(15, 17), # r_eye - r_ear
|
|
(0, 16), # nose - l_eye
|
|
(16, 18), # l_eye - l_ear
|
|
(2, 17), # r_shoulder - r_ear
|
|
(5, 18), # l_shoulder - l_ear
|
|
]
|
|
|
|
# Hand edges (0-indexed, 21 keypoints per hand)
|
|
HAND_EDGES = [
|
|
(0, 1), (1, 2), (2, 3), (3, 4), (0, 5), (5, 6), (6, 7), (7, 8),
|
|
(0, 9), (9, 10), (10, 11), (11, 12), (0, 13), (13, 14), (14, 15), (15, 16),
|
|
(0, 17), (17, 18), (18, 19), (19, 20),
|
|
]
|
|
|
|
# Colors for body limbs (BGR for OpenCV)
|
|
BODY_COLORS = [
|
|
(0, 0, 255), (0, 85, 255), (0, 170, 255), (0, 255, 255), (0, 255, 170), (0, 255, 85),
|
|
(0, 255, 0), (85, 255, 0), (170, 255, 0), (255, 255, 0), (255, 170, 0), (255, 85, 0),
|
|
(255, 0, 0), (255, 0, 85), (255, 0, 170), (255, 0, 255), (170, 0, 255), (85, 0, 255),
|
|
]
|
|
|
|
|
|
def _extract_keypoint_array(flat: list[float] | None) -> list[tuple[float, float, float]]:
|
|
"""Convert flat [x,y,c, x,y,c, ...] to list of (x, y, conf)."""
|
|
if not flat or len(flat) < 3:
|
|
return []
|
|
points = []
|
|
for i in range(0, len(flat), 3):
|
|
if i + 2 < len(flat):
|
|
x, y, c = flat[i], flat[i + 1], flat[i + 2]
|
|
points.append((float(x), float(y), float(c)))
|
|
return points
|
|
|
|
|
|
def _body_from_numpy(arr: np.ndarray) -> list[tuple[float, float, float]]:
|
|
"""Convert body array (N,2) or (N,3) to list of (x,y,conf)."""
|
|
if arr is None or arr.size == 0:
|
|
return []
|
|
arr = np.asarray(arr)
|
|
if arr.ndim != 2:
|
|
return []
|
|
points = []
|
|
for i in range(arr.shape[0]):
|
|
if arr.shape[1] >= 3:
|
|
points.append((float(arr[i, 0]), float(arr[i, 1]), float(arr[i, 2])))
|
|
elif arr.shape[1] >= 2:
|
|
points.append((float(arr[i, 0]), float(arr[i, 1]), 1.0))
|
|
else:
|
|
break
|
|
return points
|
|
|
|
|
|
def parse_pose_keypoint(pose_keypoint: Any) -> dict[str, list[tuple[float, float, float]]] | None:
|
|
"""
|
|
Parse POSE_KEYPOINT from ComfyUI (OpenPose/DWPose format) into normalized structure.
|
|
Returns dict with keys: body, hand_left, hand_right, face; each value is list of (x,y,conf).
|
|
Returns None if invalid or empty.
|
|
"""
|
|
if pose_keypoint is None:
|
|
return None
|
|
|
|
# POSE_KEYPOINT can be a list (one per image) or single dict
|
|
data = pose_keypoint
|
|
if isinstance(pose_keypoint, list):
|
|
if not pose_keypoint:
|
|
return None
|
|
data = pose_keypoint[0]
|
|
|
|
if not isinstance(data, dict):
|
|
return None
|
|
|
|
# Structure 1: {"people": [{"pose_keypoints_2d": [...], "hand_left_keypoints_2d": [...], ...}]}
|
|
people = data.get("people", [])
|
|
if people:
|
|
person = people[0]
|
|
body = _extract_keypoint_array(person.get("pose_keypoints_2d"))
|
|
hand_left = _extract_keypoint_array(person.get("hand_left_keypoints_2d"))
|
|
hand_right = _extract_keypoint_array(person.get("hand_right_keypoints_2d"))
|
|
face = _extract_keypoint_array(person.get("face_keypoints_2d"))
|
|
if body:
|
|
return {"body": body, "hand_left": hand_left, "hand_right": hand_right, "face": face}
|
|
|
|
# Structure 2: {"person_0": {"body": np.array, "hands": np.array, "face": np.array}}
|
|
for key in ("person_0", "person"):
|
|
if key in data:
|
|
person = data[key]
|
|
body = person.get("body")
|
|
if body is not None:
|
|
body = _body_from_numpy(np.asarray(body))
|
|
else:
|
|
body = []
|
|
hands = person.get("hands", np.zeros((2, 21, 2)))
|
|
hands = np.asarray(hands)
|
|
hand_left = _body_from_numpy(hands[0]) if hands.shape[0] > 0 else []
|
|
hand_right = _body_from_numpy(hands[1]) if hands.shape[0] > 1 else []
|
|
face = _body_from_numpy(np.asarray(person.get("face", [])))
|
|
if body:
|
|
return {"body": body, "hand_left": hand_left, "hand_right": hand_right, "face": face}
|
|
|
|
return None
|
|
|
|
|
|
def extract_keypoints_from_image(image: np.ndarray) -> dict[str, list[tuple[float, float, float]]] | None:
|
|
"""
|
|
Run DWPose on image to extract keypoints. Requires comfyui-controlnet-aux.
|
|
Returns parsed keypoint dict or None on failure.
|
|
"""
|
|
try:
|
|
from custom_controlnet_aux.dwpose import DwposeDetector
|
|
import comfy.model_management as model_management
|
|
except ImportError:
|
|
return None
|
|
|
|
try:
|
|
model = DwposeDetector.from_pretrained(
|
|
"yzd-v/DWPose",
|
|
"yzd-v/DWPose",
|
|
det_filename="yolox_l.onnx",
|
|
pose_filename="dw-ll_ucoco_384.onnx",
|
|
torchscript_device=model_management.get_torch_device(),
|
|
)
|
|
# DWPose expects HWC uint8; annotators often expect BGR
|
|
if image.dtype != np.uint8:
|
|
image = (np.clip(image, 0, 1) * 255).astype(np.uint8)
|
|
if len(image.shape) == 2:
|
|
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
|
elif image.shape[2] == 3:
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
pose_img, openpose_dict = model(
|
|
image,
|
|
include_hand=True,
|
|
include_face=True,
|
|
include_body=True,
|
|
detect_resolution=min(image.shape[0], image.shape[1], 512),
|
|
)
|
|
del model
|
|
return parse_pose_keypoint(openpose_dict)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def compute_torso_pivot(keypoints: dict[str, list[tuple[float, float, float]]]) -> tuple[float, float] | None:
|
|
"""
|
|
Compute torso center from body keypoints. Uses neck, shoulders, mid_hip, hips.
|
|
Returns (cx, cy) or None if insufficient keypoints.
|
|
"""
|
|
body = keypoints.get("body", [])
|
|
if len(body) < 2:
|
|
return None
|
|
|
|
xs, ys, total_conf = 0.0, 0.0, 0.0
|
|
for idx in TORSO_INDICES:
|
|
if idx >= len(body):
|
|
continue
|
|
x, y, c = body[idx]
|
|
if c > 0:
|
|
xs += x * c
|
|
ys += y * c
|
|
total_conf += c
|
|
|
|
if total_conf <= 0:
|
|
return None
|
|
|
|
return (xs / total_conf, ys / total_conf)
|
|
|
|
|
|
def _compute_adaptive_depth_scale(
|
|
keypoints: dict[str, list[tuple[float, float, float]]],
|
|
pivot: tuple[float, float],
|
|
) -> float:
|
|
"""
|
|
Compute adaptive depth scale from body proportions (shoulder width).
|
|
Normalizes to typical pose so rotation looks consistent across image sizes.
|
|
"""
|
|
body = keypoints.get("body", [])
|
|
if len(body) < 6:
|
|
return 0.4
|
|
x2, _, c2 = body[2] # r_shoulder
|
|
x5, _, c5 = body[5] # l_shoulder
|
|
if c2 <= 0 or c5 <= 0:
|
|
return 0.4
|
|
shoulder_width = abs(x2 - x5)
|
|
if shoulder_width < 1:
|
|
return 0.4
|
|
# Typical shoulder width ~100-150px in normalized poses; scale so base 0.4 holds
|
|
return 0.4 * (120.0 / shoulder_width)
|
|
|
|
|
|
def rotate_keypoints_3d(
|
|
keypoints: dict[str, list[tuple[float, float, float]]],
|
|
pivot: tuple[float, float],
|
|
degrees: float,
|
|
direction: str,
|
|
depth_scale: float | None = None,
|
|
) -> tuple[dict[str, list[tuple[float, float, float]]], dict[str, list[float]]]:
|
|
"""
|
|
Rotate all keypoints around Y-axis through pivot. Left = positive angle (CCW), right = negative.
|
|
Uses limb-specific depth inference per OpenPose COCO body indices.
|
|
Returns (rotated_keypoints, depth_per_point) for occlusion and depth-ordered rendering.
|
|
"""
|
|
theta_deg = degrees if direction == "left" else -degrees
|
|
theta = math.radians(theta_deg)
|
|
cx, cy = pivot
|
|
|
|
if depth_scale is None:
|
|
depth_scale = _compute_adaptive_depth_scale(keypoints, pivot)
|
|
|
|
cos_t = math.cos(theta)
|
|
sin_t = math.sin(theta)
|
|
|
|
def rotate_point(
|
|
x: float, y: float, scale: float
|
|
) -> tuple[float, float, float]:
|
|
# Side-aware: (x - cx) gives correct sign for left/right body sides
|
|
z = (x - cx) * scale
|
|
x_new = (x - cx) * cos_t - z * sin_t + cx
|
|
y_new = y
|
|
z_new = (x - cx) * sin_t + z * cos_t # depth after rotation
|
|
return (x_new, y_new, z_new)
|
|
|
|
result: dict[str, list[tuple[float, float, float]]] = {}
|
|
depths: dict[str, list[float]] = {}
|
|
|
|
for part, points in keypoints.items():
|
|
rotated: list[tuple[float, float, float]] = []
|
|
part_depths: list[float] = []
|
|
|
|
for i, (x, y, c) in enumerate(points):
|
|
if c <= 0:
|
|
rotated.append((x, y, c))
|
|
part_depths.append(0.0)
|
|
continue
|
|
|
|
scale = BODY_DEPTH_SCALES.get(i, HAND_DEPTH_SCALE) if part == "body" else HAND_DEPTH_SCALE
|
|
nx, ny, z_new = rotate_point(x, y, depth_scale * scale)
|
|
|
|
# Occlusion: hide points far behind torso
|
|
if z_new < OCCLUSION_DEPTH_THRESHOLD:
|
|
rotated.append((nx, ny, 0.0))
|
|
else:
|
|
rotated.append((nx, ny, c))
|
|
part_depths.append(z_new)
|
|
|
|
result[part] = rotated
|
|
depths[part] = part_depths
|
|
|
|
# Face visibility: when rotated >90°, hide face/head keypoints (back of head)
|
|
if abs(theta_deg) >= 90:
|
|
if "body" in result:
|
|
body = result["body"]
|
|
for idx in FACE_INDICES:
|
|
if idx < len(body):
|
|
x, y, c = body[idx]
|
|
body[idx] = (x, y, 0.0)
|
|
# Also hide separate face keypoints (OpenPose face mesh)
|
|
if "face" in result:
|
|
result["face"] = [(x, y, 0.0) for x, y, c in result["face"]]
|
|
|
|
return result, depths
|
|
|
|
|
|
def join_broken_segments(
|
|
keypoints: dict[str, list[tuple[float, float, float]]],
|
|
threshold: float,
|
|
) -> dict[str, list[tuple[float, float, float]]]:
|
|
"""
|
|
Optional: join broken segments when one endpoint visible, other missing.
|
|
Currently a no-op (disabled) per plan.
|
|
"""
|
|
return keypoints
|
|
|
|
|
|
def render_openpose_image(
|
|
keypoints: dict[str, list[tuple[float, float, float]]],
|
|
width: int,
|
|
height: int,
|
|
depths: dict[str, list[float]] | None = None,
|
|
) -> np.ndarray:
|
|
"""
|
|
Draw OpenPose skeleton on blank canvas. White background, colored limbs and joints.
|
|
When depths is provided, limbs are drawn back-to-front for correct occlusion.
|
|
"""
|
|
canvas = np.ones((height, width, 3), dtype=np.uint8) * 255
|
|
|
|
body = keypoints.get("body", [])
|
|
body_depths = depths.get("body", []) if depths else []
|
|
|
|
# Build limb list with depth for sorting (mean depth of endpoints)
|
|
limb_data: list[tuple[int, int, int, float]] = []
|
|
for i, (a, b) in enumerate(BODY_LIMBS):
|
|
if a >= len(body) or b >= len(body):
|
|
continue
|
|
x1, y1, c1 = body[a]
|
|
x2, y2, c2 = body[b]
|
|
if c1 <= 0 or c2 <= 0:
|
|
continue
|
|
depth_a = body_depths[a] if a < len(body_depths) else 0.0
|
|
depth_b = body_depths[b] if b < len(body_depths) else 0.0
|
|
mean_depth = (depth_a + depth_b) / 2
|
|
limb_data.append((a, b, i, mean_depth))
|
|
|
|
# Draw body limbs back-to-front (far limbs first)
|
|
limb_data.sort(key=lambda t: t[3])
|
|
for a, b, i, _ in limb_data:
|
|
x1, y1, c1 = body[a]
|
|
x2, y2, c2 = body[b]
|
|
color = BODY_COLORS[i % len(BODY_COLORS)]
|
|
cv2.line(canvas, (int(x1), int(y1)), (int(x2), int(y2)), color, 4)
|
|
|
|
# Draw body joints (front-facing first so they appear on top)
|
|
joint_data = [(i, x, y, c) for i, (x, y, c) in enumerate(body) if c > 0]
|
|
if body_depths:
|
|
joint_data.sort(key=lambda t: -body_depths[t[0]] if t[0] < len(body_depths) else 0)
|
|
for i, x, y, c in joint_data:
|
|
color = BODY_COLORS[i % len(BODY_COLORS)]
|
|
cv2.circle(canvas, (int(x), int(y)), 4, color, -1)
|
|
|
|
# Draw hands (with depth ordering if available)
|
|
for hand_key, hand_pts in [("hand_left", keypoints.get("hand_left", [])),
|
|
("hand_right", keypoints.get("hand_right", []))]:
|
|
if len(hand_pts) < 2:
|
|
continue
|
|
hand_depths = depths.get(hand_key, []) if depths else []
|
|
|
|
# Sort hand edges by depth for proper occlusion
|
|
edges_with_depth: list[tuple[int, int, float]] = []
|
|
for a, b in HAND_EDGES:
|
|
if a >= len(hand_pts) or b >= len(hand_pts):
|
|
continue
|
|
x1, y1, c1 = hand_pts[a]
|
|
x2, y2, c2 = hand_pts[b]
|
|
if c1 <= 0 or c2 <= 0:
|
|
continue
|
|
d_a = hand_depths[a] if a < len(hand_depths) else 0.0
|
|
d_b = hand_depths[b] if b < len(hand_depths) else 0.0
|
|
edges_with_depth.append((a, b, (d_a + d_b) / 2))
|
|
if hand_depths:
|
|
edges_with_depth.sort(key=lambda t: t[2])
|
|
for a, b, _ in edges_with_depth:
|
|
x1, y1, c1 = hand_pts[a]
|
|
x2, y2, c2 = hand_pts[b]
|
|
cv2.line(canvas, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
|
|
|
|
joint_order = list(range(len(hand_pts)))
|
|
if hand_depths and len(hand_depths) == len(hand_pts):
|
|
joint_order.sort(key=lambda idx: -hand_depths[idx] if hand_pts[idx][2] > 0 else 999)
|
|
for idx in joint_order:
|
|
x, y, c = hand_pts[idx]
|
|
if c > 0:
|
|
cv2.circle(canvas, (int(x), int(y)), 3, (0, 0, 255), -1)
|
|
|
|
return canvas
|
|
|
|
|
|
def rotate_openpose(
|
|
image: np.ndarray,
|
|
pose_keypoint: Any | None,
|
|
direction: str,
|
|
degrees: float,
|
|
) -> tuple[np.ndarray, bool]:
|
|
"""
|
|
Main pipeline: extract/parse keypoints, detect torso, rotate, render.
|
|
Returns (output_image, success). On failure, returns (input_image, False).
|
|
"""
|
|
h, w = image.shape[:2]
|
|
|
|
# Get keypoints
|
|
if pose_keypoint is not None:
|
|
keypoints = parse_pose_keypoint(pose_keypoint)
|
|
else:
|
|
keypoints = extract_keypoints_from_image(image)
|
|
|
|
if keypoints is None:
|
|
return image, False
|
|
|
|
# Torso pivot
|
|
pivot = compute_torso_pivot(keypoints)
|
|
if pivot is None:
|
|
return image, False
|
|
|
|
# Rotate (returns keypoints and depth for occlusion/ordering)
|
|
rotated, depths = rotate_keypoints_3d(keypoints, pivot, degrees, direction)
|
|
rotated = join_broken_segments(rotated, 0.0)
|
|
|
|
# Render with depth-ordered drawing
|
|
rendered = render_openpose_image(rotated, w, h, depths)
|
|
return rendered, True
|