1168 lines
46 KiB
Python
1168 lines
46 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, ControlNet / DWPose COCO-18 format)
|
||
# 0=nose, 1=neck, 2=r_shoulder, 3=r_elbow, 4=r_wrist, 5=l_shoulder, 6=l_elbow, 7=l_wrist,
|
||
# 8=r_hip, 9=r_knee, 10=r_ankle, 11=l_hip, 12=l_knee, 13=l_ankle,
|
||
# 14=r_eye, 15=l_eye, 16=r_ear, 17=l_ear
|
||
# Note: no mid_hip in COCO-18. Legs connect directly from neck to each hip.
|
||
TORSO_INDICES = [1, 2, 5, 8, 11] # neck, r_shoulder, l_shoulder, r_hip, l_hip
|
||
|
||
# Face/head keypoints (hidden when rotated >150° - back of head visible)
|
||
# Keeps face visible when facing to the side (~90°) so at least one eye remains
|
||
FACE_INDICES = {0, 14, 15, 16, 17} # nose, r_eye, l_eye, r_ear, l_ear (COCO-18)
|
||
FACE_HIDE_ANGLE = 150 # degrees - hide face only when head is turned away
|
||
|
||
# Limb-specific depth scales (ControlNet COCO-18 body indices 0-17)
|
||
# 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
|
||
1: 0.20, # neck
|
||
2: 0.35, # r_shoulder
|
||
3: 0.50, # r_elbow
|
||
4: 0.55, # r_wrist
|
||
5: 0.35, # l_shoulder
|
||
6: 0.50, # l_elbow
|
||
7: 0.55, # l_wrist
|
||
8: 0.30, # r_hip (COCO-18)
|
||
9: 0.45, # r_knee
|
||
10: 0.50, # r_ankle
|
||
11: 0.30, # l_hip (COCO-18)
|
||
12: 0.45, # l_knee
|
||
13: 0.50, # l_ankle
|
||
14: 0.28, # r_eye (COCO-18)
|
||
15: 0.28, # l_eye
|
||
16: 0.30, # r_ear
|
||
17: 0.30, # l_ear
|
||
}
|
||
|
||
# Default depth scale for hands (21 keypoints each - no per-index in OpenPose hand spec)
|
||
HAND_DEPTH_SCALE = 0.45
|
||
|
||
# Face mesh (70 keypoints) - same region as head; use head-like depth so face rotates with skull
|
||
FACE_DEPTH_SCALE = 0.30
|
||
|
||
# Static limb definitions: (a, b) connection, draw layer (0=head first, 4=feet last), color index
|
||
# Uses ControlNet COCO-18 indices. No mid_hip — legs connect neck→hip directly.
|
||
BODY_LIMBS = [
|
||
((0, 1), 0, 0), # nose - neck
|
||
((0, 14), 0, 0), # nose - r_eye (COCO-18: 14=r_eye)
|
||
((0, 15), 0, 0), # nose - l_eye (COCO-18: 15=l_eye)
|
||
((14, 16), 0, 0), # r_eye - r_ear (COCO-18: 16=r_ear)
|
||
((15, 17), 0, 0), # l_eye - l_ear (COCO-18: 17=l_ear)
|
||
((1, 2), 0, 0), # neck - r_shoulder
|
||
((1, 5), 0, 0), # neck - l_shoulder
|
||
((2, 3), 1, 2), # r_shoulder - r_elbow
|
||
((3, 4), 2, 3), # r_elbow - r_wrist
|
||
((5, 6), 1, 4), # l_shoulder - l_elbow
|
||
((6, 7), 2, 5), # l_elbow - l_wrist
|
||
((1, 8), 0, 6), # neck - r_hip (spine right, COCO-18: 8=r_hip)
|
||
((1, 11), 0, 6), # neck - l_hip (spine left, COCO-18: 11=l_hip)
|
||
((8, 9), 3, 7), # r_hip - r_knee
|
||
((9, 10), 4, 8), # r_knee - r_ankle
|
||
((11, 12), 3, 9), # l_hip - l_knee
|
||
((12, 13), 4, 10),# l_knee - l_ankle
|
||
]
|
||
|
||
# ControlNet OpenPose standard colors (BGR) - matches expected figure appearance
|
||
# Red chest -> orange/yellow arms -> green spine -> teal/blue/purple legs
|
||
BODY_COLORS = [
|
||
(0, 0, 255), # red - neck-shoulders (chest line)
|
||
(0, 0, 255),
|
||
(0, 85, 255), # orange - r upper arm
|
||
(0, 170, 255), # yellow - r forearm
|
||
(0, 255, 85), # green - l upper arm
|
||
(85, 255, 0), # lime green - l forearm
|
||
(0, 255, 0), # green - spine (neck to mid_hip)
|
||
(0, 255, 170), # teal - r thigh
|
||
(0, 255, 255), # cyan - r knee
|
||
(170, 255, 0), # light blue - r shin
|
||
(255, 0, 0), # blue - l thigh
|
||
(255, 85, 0), # dark blue - l knee
|
||
(255, 0, 170), # purple - l shin
|
||
]
|
||
|
||
# Static joint definitions: index -> (draw layer, color index) — ControlNet COCO-18
|
||
JOINT_LAYERS = {
|
||
0: (0, 0), # nose
|
||
1: (0, 0), # neck
|
||
2: (1, 2), # r_shoulder
|
||
3: (1, 2), # r_elbow
|
||
4: (1, 3), # r_wrist
|
||
5: (1, 4), # l_shoulder
|
||
6: (1, 4), # l_elbow
|
||
7: (1, 5), # l_wrist
|
||
8: (2, 7), # r_hip (COCO-18)
|
||
9: (2, 8), # r_knee
|
||
10: (3, 8), # r_ankle
|
||
11: (2, 9), # l_hip (COCO-18)
|
||
12: (2, 10), # l_knee
|
||
13: (3, 10), # l_ankle
|
||
14: (0, 0), # r_eye (COCO-18)
|
||
15: (0, 0), # l_eye
|
||
16: (0, 0), # r_ear
|
||
17: (0, 0), # l_ear
|
||
}
|
||
|
||
# Valid keypoint indices for limb drawing (COCO-18: indices 0-17)
|
||
MAIN_BODY_INDICES = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17}
|
||
|
||
# Hand edges (0-indexed, 21 keypoints per hand) - used when hands enabled
|
||
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),
|
||
]
|
||
|
||
# Face mesh edges (0-indexed, 70 keypoints) - OpenPose FACE_PAIRS_RENDER from faceParameters.hpp
|
||
FACE_EDGES = [
|
||
(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7), (7, 8), (8, 9), (9, 10),
|
||
(10, 11), (11, 12), (12, 13), (13, 14), (14, 15), (15, 16),
|
||
(17, 18), (18, 19), (19, 20), (20, 21),
|
||
(22, 23), (23, 24), (24, 25), (25, 26),
|
||
(27, 28), (28, 29), (29, 30),
|
||
(31, 32), (32, 33), (33, 34), (34, 35),
|
||
(36, 37), (37, 38), (38, 39), (39, 40), (40, 41), (41, 36),
|
||
(42, 43), (43, 44), (44, 45), (45, 46), (46, 47), (47, 42),
|
||
(48, 49), (49, 50), (50, 51), (51, 52), (52, 53), (53, 54), (54, 55), (55, 56),
|
||
(56, 57), (57, 58), (58, 59), (59, 48),
|
||
(60, 61), (61, 62), (62, 63), (63, 64), (64, 65), (65, 66), (66, 67), (67, 60),
|
||
]
|
||
|
||
|
||
def _keypoints_to_openpose_dict(
|
||
keypoints: dict[str, list[tuple[float, float, float]]],
|
||
) -> dict[str, list[float]]:
|
||
"""Convert internal keypoint dict to OpenPose/ComfyUI format (flat x,y,c arrays)."""
|
||
body = keypoints.get("body", [])
|
||
hand_left = keypoints.get("hand_left", [])
|
||
hand_right = keypoints.get("hand_right", [])
|
||
face = keypoints.get("face", [])
|
||
|
||
def flatten(pts: list[tuple[float, float, float]]) -> list[float]:
|
||
result = []
|
||
for x, y, c in pts:
|
||
result.extend([float(x), float(y), float(c)])
|
||
return result
|
||
|
||
person = {
|
||
"pose_keypoints_2d": flatten(body),
|
||
"hand_left_keypoints_2d": flatten(hand_left),
|
||
"hand_right_keypoints_2d": flatten(hand_right),
|
||
"face_keypoints_2d": flatten(face),
|
||
}
|
||
return {"people": [person]}
|
||
|
||
|
||
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 _scale_keypoints_to_image(
|
||
keypoints: dict[str, list[tuple[float, float, float]]],
|
||
width: int,
|
||
height: int,
|
||
) -> dict[str, list[tuple[float, float, float]]]:
|
||
"""
|
||
Scale keypoints to image dimensions.
|
||
- Normalized [0,1] coords: scale by width, height
|
||
- Pixel coords from smaller resolution (e.g. DWPose at 384): scale to match image
|
||
"""
|
||
all_x, all_y = [], []
|
||
for pts in keypoints.values():
|
||
for x, y, c in pts:
|
||
if c > 0:
|
||
all_x.append(x)
|
||
all_y.append(y)
|
||
|
||
if not all_x or not all_y:
|
||
return keypoints
|
||
|
||
max_x, max_y = max(all_x), max(all_y)
|
||
min_x, min_y = min(all_x), min(all_y)
|
||
|
||
# Normalized coords in [0,1] (ComfyUI POSE_KEYPOINT may use this)
|
||
if max_x <= 1.5 and max_y <= 1.5 and min_x >= -0.1 and min_y >= -0.1:
|
||
scale_x, scale_y = width, height
|
||
# Pixel coords from smaller resolution (e.g. DWPose at 384)
|
||
# Only scale when max is clearly from a smaller resolution; avoid scaling when
|
||
# the figure is already at image size but doesn't fill the frame (causes ~40% blow-up)
|
||
elif max_x < width * 0.76 and max_y < height * 0.76:
|
||
scale_x = width / max(max_x, 1)
|
||
scale_y = height / max(max_y, 1)
|
||
if abs(scale_x - 1.0) < 0.01 and abs(scale_y - 1.0) < 0.01:
|
||
return keypoints
|
||
else:
|
||
return keypoints
|
||
|
||
result: dict[str, list[tuple[float, float, float]]] = {}
|
||
for part, pts in keypoints.items():
|
||
scaled = [
|
||
(x * scale_x, y * scale_y, c) if c > 0 else (x, y, c)
|
||
for x, y, c in pts
|
||
]
|
||
result[part] = scaled
|
||
return result
|
||
|
||
|
||
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}
|
||
|
||
# Structure 3: flat {"pose_keypoints_2d": [...], ...} at top level (no "people")
|
||
body = _extract_keypoint_array(data.get("pose_keypoints_2d"))
|
||
if body:
|
||
hand_left = _extract_keypoint_array(data.get("hand_left_keypoints_2d"))
|
||
hand_right = _extract_keypoint_array(data.get("hand_right_keypoints_2d"))
|
||
face = _extract_keypoint_array(data.get("face_keypoints_2d"))
|
||
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.
|
||
Supports "no torso segment" variants: derives neck from shoulders and mid_hip from hips
|
||
when those keypoints are missing (c<=0) but shoulders/hips are present.
|
||
Returns (cx, cy) or None if insufficient keypoints.
|
||
"""
|
||
body = keypoints.get("body", [])
|
||
if len(body) < 2:
|
||
return None
|
||
|
||
def get_pt(idx: int) -> tuple[float, float, float] | None:
|
||
if idx >= len(body):
|
||
return None
|
||
x, y, c = body[idx]
|
||
return (x, y, c) if c > 0 else None
|
||
|
||
xs, ys, total_conf = 0.0, 0.0, 0.0
|
||
|
||
for idx in TORSO_INDICES:
|
||
pt = get_pt(idx)
|
||
if pt is not None:
|
||
x, y, c = pt
|
||
xs += x * c
|
||
ys += y * c
|
||
total_conf += c
|
||
|
||
# Fallback: derive neck from shoulders if missing
|
||
# COCO-18 has no mid_hip; both hips (8=r_hip, 11=l_hip) are already in TORSO_INDICES
|
||
if get_pt(1) is None:
|
||
r_sh = get_pt(2)
|
||
l_sh = get_pt(5)
|
||
if r_sh is not None and l_sh is not None:
|
||
nx = (r_sh[0] + l_sh[0]) / 2
|
||
ny = (r_sh[1] + l_sh[1]) / 2
|
||
xs += nx
|
||
ys += ny
|
||
total_conf += 1.0
|
||
|
||
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)
|
||
|
||
|
||
# Perspective scale_y: bridges the raw (cy - y) pixel value to z-units.
|
||
# In advanced mode the coefficient is derived geometrically from ADVANCED_CAMERA_ELEVATION_DEG;
|
||
# this constant keeps the formula consistent: z += perspective * (cy - y) * PERSPECTIVE_SCALE_Y.
|
||
PERSPECTIVE_SCALE_Y = 0.3
|
||
|
||
# Advanced mode: fixed physical camera / figure assumptions.
|
||
# Camera is elevated 35° above horizontal, looking at a standing figure that fills
|
||
# roughly 80-100% (average ~90%) of the image height.
|
||
ADVANCED_CAMERA_ELEVATION_DEG = 35.0
|
||
ADVANCED_FIGURE_FILL = 0.90
|
||
|
||
# Ankle y-parity threshold: fraction of figure height (neck→lower ankle) within which both ankles
|
||
# are considered to be at the same elevation (standing upright on flat ground).
|
||
UPRIGHT_ANKLE_THRESHOLD = 0.12
|
||
|
||
|
||
def _detect_upright_figure(
|
||
keypoints: dict[str, list[tuple[float, float, float]]],
|
||
) -> tuple[bool, tuple[float, float, float] | None, tuple[float, float, float] | None]:
|
||
"""
|
||
Return (is_upright, neck_pt, ankle_mid_pt).
|
||
Upright is True when neck (1) and both ankles (11=r, 14=l) are detected and both ankle y values
|
||
differ by less than UPRIGHT_ANKLE_THRESHOLD * figure_height (neck y to lower ankle y).
|
||
"""
|
||
body = keypoints.get("body", [])
|
||
|
||
def get_pt(idx: int) -> tuple[float, float, float] | None:
|
||
if idx >= len(body):
|
||
return None
|
||
x, y, c = body[idx]
|
||
return (x, y, c) if c > 0 else None
|
||
|
||
neck = get_pt(1)
|
||
r_ankle = get_pt(10) # COCO-18: r_ankle=10
|
||
l_ankle = get_pt(13) # COCO-18: l_ankle=13
|
||
|
||
if neck is None or r_ankle is None or l_ankle is None:
|
||
return False, None, None
|
||
|
||
lower_ankle_y = max(r_ankle[1], l_ankle[1])
|
||
figure_height = abs(lower_ankle_y - neck[1])
|
||
if figure_height < 1:
|
||
return False, None, None
|
||
|
||
ankle_y_diff = abs(r_ankle[1] - l_ankle[1])
|
||
if ankle_y_diff > UPRIGHT_ANKLE_THRESHOLD * figure_height:
|
||
return False, None, None
|
||
|
||
ankle_mid: tuple[float, float, float] = (
|
||
(r_ankle[0] + l_ankle[0]) / 2,
|
||
(r_ankle[1] + l_ankle[1]) / 2,
|
||
1.0,
|
||
)
|
||
return True, neck, ankle_mid
|
||
|
||
|
||
def _compute_advanced_params(image_height: int) -> tuple[float, float]:
|
||
"""
|
||
Derive perspective and focal_length for advanced mode from fixed camera/figure assumptions.
|
||
|
||
Camera elevation θ = ADVANCED_CAMERA_ELEVATION_DEG (35°) above horizontal.
|
||
Figure fills ADVANCED_FIGURE_FILL (~90%) of image height.
|
||
|
||
perspective: maps vertical pixel distance from pivot to depth.
|
||
Formula: z += perspective * (cy - y) * PERSPECTIVE_SCALE_Y
|
||
For geometric accuracy: perspective * PERSPECTIVE_SCALE_Y = tan(θ)
|
||
→ perspective = tan(35°) / 0.3 ≈ 2.33
|
||
|
||
focal_length: perspective projection strength, scaled to image resolution.
|
||
From pinhole geometry: focal_length = image_height * fill / tan(θ)
|
||
At 35° with 90% fill → focal_length ≈ 1.29 × image_height
|
||
"""
|
||
tan_elev = math.tan(math.radians(ADVANCED_CAMERA_ELEVATION_DEG))
|
||
perspective = tan_elev / PERSPECTIVE_SCALE_Y
|
||
focal_length = image_height * ADVANCED_FIGURE_FILL / tan_elev
|
||
return perspective, focal_length
|
||
|
||
|
||
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,
|
||
mode: str = "simple",
|
||
perspective: float = 0.0,
|
||
focal_length: float = 800.0,
|
||
) -> dict[str, list[tuple[float, float, float]]]:
|
||
"""
|
||
Rotate all keypoints around Y-axis through pivot. Counterclockwise = positive angle, clockwise = negative.
|
||
Uses limb-specific depth inference per OpenPose COCO body indices for 3D-style projection.
|
||
|
||
mode: "simple" = orthographic output with perspective-modulated depth; "advanced" = perspective projection.
|
||
perspective: 0 = eye level, positive = camera above (isometric), negative = camera below.
|
||
focal_length: (advanced only) controls perspective strength; higher = flatter, lower = more foreshortening.
|
||
|
||
In advanced mode, when the figure is detected as standing upright (neck + matching-y ankles),
|
||
the rotation axis is the 3D spine computed from the ankle midpoint through the neck rather than
|
||
a pure vertical Y-axis. This keeps the neck stationary and tilts the axis to match the figure's lean.
|
||
"""
|
||
theta_deg = degrees if direction == "counterclockwise" 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)
|
||
use_perspective_projection = mode == "advanced"
|
||
|
||
# --- Advanced mode: detect upright figure and build spine axis ---
|
||
# axis_point and axis_dir are in pivot-relative 3D space (origin = pivot).
|
||
axis_point: np.ndarray | None = None
|
||
axis_dir: np.ndarray | None = None
|
||
|
||
# Upright spine axis applies in both modes — neck is always the rotation anchor when conditions met
|
||
is_upright, neck_pt, ankle_mid_pt = _detect_upright_figure(keypoints)
|
||
if is_upright and neck_pt is not None and ankle_mid_pt is not None:
|
||
def _infer_z(px: float, py: float, limb_scale: float) -> float:
|
||
return (px - cx) * depth_scale * limb_scale + perspective * (cy - py) * PERSPECTIVE_SCALE_Y
|
||
|
||
# Neck 3D (pivot-relative)
|
||
nz = _infer_z(neck_pt[0], neck_pt[1], BODY_DEPTH_SCALES.get(1, 0.20))
|
||
neck_3d = np.array([neck_pt[0] - cx, neck_pt[1] - cy, nz])
|
||
|
||
# Ankle midpoint 3D (pivot-relative) — average depth of both ankles (COCO-18: 10/13)
|
||
r_scale = BODY_DEPTH_SCALES.get(10, 0.50)
|
||
l_scale = BODY_DEPTH_SCALES.get(13, 0.50)
|
||
body = keypoints.get("body", [])
|
||
r_ankle = body[10]
|
||
l_ankle = body[13]
|
||
z_r = _infer_z(r_ankle[0], r_ankle[1], r_scale)
|
||
z_l = _infer_z(l_ankle[0], l_ankle[1], l_scale)
|
||
ankle_3d = np.array([
|
||
ankle_mid_pt[0] - cx,
|
||
ankle_mid_pt[1] - cy,
|
||
(z_r + z_l) / 2,
|
||
])
|
||
|
||
spine = neck_3d - ankle_3d
|
||
spine_norm = float(np.linalg.norm(spine))
|
||
if spine_norm > 1e-6:
|
||
axis_dir = spine / spine_norm
|
||
axis_point = neck_3d # axis passes through neck
|
||
|
||
def _infer_z_point(x: float, y: float, scale: float) -> float:
|
||
# scale is depth_scale * limb_scale
|
||
return (x - cx) * scale + perspective * (cy - y) * PERSPECTIVE_SCALE_Y
|
||
|
||
def rotate_point(
|
||
x: float, y: float, scale: float
|
||
) -> tuple[float, float, float]:
|
||
z = _infer_z_point(x, y, scale)
|
||
|
||
if axis_dir is not None and axis_point is not None:
|
||
# Rodrigues rotation around the spine axis (advanced + upright)
|
||
p = np.array([x - cx, y - cy, z])
|
||
p_rel = p - axis_point
|
||
k = axis_dir
|
||
p_rot = (p_rel * cos_t
|
||
+ np.cross(k, p_rel) * sin_t
|
||
+ k * float(np.dot(k, p_rel)) * (1.0 - cos_t))
|
||
p_world = p_rot + axis_point
|
||
x_rot, y_rot_3d, z_rot = float(p_world[0]), float(p_world[1]), float(p_world[2])
|
||
if use_perspective_projection:
|
||
z_cam = focal_length + z_rot
|
||
z_cam = max(z_cam, 0.1 * focal_length)
|
||
# Perspective division applies only to X (horizontal foreshortening during yaw).
|
||
# Y is kept orthographic: dividing by z_cam would vertically stretch the figure
|
||
# because the elevated camera assigns large z values to high/low keypoints,
|
||
# causing them to shift far outside the image even at near-zero rotation angles.
|
||
return (cx + focal_length * x_rot / z_cam,
|
||
cy + y_rot_3d,
|
||
z_rot)
|
||
return (x_rot + cx, y_rot_3d + cy, z_rot)
|
||
|
||
# Default: Y-axis rotation in XZ plane (simple mode or non-upright advanced)
|
||
x_rot = (x - cx) * cos_t - z * sin_t
|
||
z_rot = (x - cx) * sin_t + z * cos_t
|
||
if use_perspective_projection:
|
||
z_cam = focal_length + z_rot
|
||
z_cam = max(z_cam, 0.1 * focal_length)
|
||
return (cx + focal_length * x_rot / z_cam,
|
||
y, # Y unchanged: pure yaw does not affect vertical screen position
|
||
z_rot)
|
||
return (x_rot + cx, y, z_rot)
|
||
|
||
result: dict[str, list[tuple[float, float, float]]] = {}
|
||
|
||
for part, points in keypoints.items():
|
||
rotated: list[tuple[float, float, float]] = []
|
||
|
||
for i, (x, y, c) in enumerate(points):
|
||
if c <= 0:
|
||
rotated.append((x, y, c))
|
||
continue
|
||
|
||
if part == "body":
|
||
scale = BODY_DEPTH_SCALES.get(i, HAND_DEPTH_SCALE)
|
||
elif part == "face":
|
||
scale = FACE_DEPTH_SCALE
|
||
else:
|
||
scale = HAND_DEPTH_SCALE
|
||
nx, ny, _ = rotate_point(x, y, depth_scale * scale)
|
||
rotated.append((nx, ny, c))
|
||
|
||
result[part] = rotated
|
||
|
||
# Face visibility — two regimes:
|
||
#
|
||
# Forward-facing (both eyes detected in original pose):
|
||
# Hemisphere model — a face keypoint whose original normalised lateral position is
|
||
# t = (x − face_cx) / half_w is hidden once sign_dir × t > cos(|θ|).
|
||
# This is geometrically exact for a hemispherical head: edge features (ears) disappear
|
||
# first; the nose disappears last (only past 90°).
|
||
# sign_dir = +1 for CCW (right-of-image side recedes), −1 for CW (left side recedes).
|
||
#
|
||
# Non-forward-facing (one or both eyes absent):
|
||
# Original binary threshold — hide all face points once |θ| ≥ FACE_HIDE_ANGLE.
|
||
|
||
if abs(theta_deg) > 0:
|
||
orig_body = keypoints.get("body", [])
|
||
orig_face = keypoints.get("face", [])
|
||
|
||
r_eye_orig = orig_body[14] if len(orig_body) > 14 else None
|
||
l_eye_orig = orig_body[15] if len(orig_body) > 15 else None
|
||
both_eyes = (
|
||
r_eye_orig is not None and r_eye_orig[2] > 0
|
||
and l_eye_orig is not None and l_eye_orig[2] > 0
|
||
)
|
||
|
||
if both_eyes:
|
||
# Face centre and half-width from original pose.
|
||
# Prefer the face mesh (covers full jaw); fall back to eye spacing.
|
||
face_xs = [x for x, y, c in orig_face if c > 0] if orig_face else []
|
||
if len(face_xs) >= 2:
|
||
face_cx_orig = (min(face_xs) + max(face_xs)) / 2.0
|
||
half_w = (max(face_xs) - min(face_xs)) / 2.0
|
||
else:
|
||
# Eyes span roughly the middle third of the face, so half_w ≈ 1.5 × eye_sep
|
||
face_cx_orig = (r_eye_orig[0] + l_eye_orig[0]) / 2.0
|
||
half_w = abs(l_eye_orig[0] - r_eye_orig[0]) * 1.5
|
||
half_w = max(half_w, 1.0)
|
||
|
||
cos_theta = math.cos(math.radians(abs(theta_deg)))
|
||
sign_dir = 1.0 if theta_deg > 0 else -1.0 # +1 CCW, −1 CW
|
||
|
||
# Hide face mesh keypoints (70-point OpenPose face)
|
||
if "face" in result and orig_face:
|
||
face_out = result["face"]
|
||
for i, (x_orig, y_orig, c_orig) in enumerate(orig_face):
|
||
if c_orig <= 0 or i >= len(face_out):
|
||
continue
|
||
t = (x_orig - face_cx_orig) / half_w
|
||
if sign_dir * t > cos_theta:
|
||
px, py, _ = face_out[i]
|
||
face_out[i] = (px, py, 0.0)
|
||
|
||
# Hide body face keypoints: nose(0), r_eye(14), l_eye(15), r_ear(16), l_ear(17)
|
||
if "body" in result:
|
||
body_out = result["body"]
|
||
for idx in FACE_INDICES:
|
||
if idx >= len(orig_body) or idx >= len(body_out):
|
||
continue
|
||
ox, oy, oc = orig_body[idx]
|
||
if oc <= 0:
|
||
continue
|
||
t = (ox - face_cx_orig) / half_w
|
||
if sign_dir * t > cos_theta:
|
||
px, py, _ = body_out[idx]
|
||
body_out[idx] = (px, py, 0.0)
|
||
|
||
else:
|
||
# Non-forward-facing: original binary threshold
|
||
if abs(theta_deg) >= FACE_HIDE_ANGLE:
|
||
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)
|
||
if "face" in result:
|
||
result["face"] = [(x, y, 0.0) for x, y, c in result["face"]]
|
||
|
||
return result
|
||
|
||
|
||
def _get_anchor_point(keypoints: dict[str, list[tuple[float, float, float]]]) -> tuple[float, float] | None:
|
||
"""
|
||
Get anchor point from body keypoints for positioning.
|
||
Prefer neck (1); if missing, use midpoint of shoulders (2, 5); if only one shoulder, use it.
|
||
Returns (x, y) or None if insufficient keypoints.
|
||
"""
|
||
body = keypoints.get("body", [])
|
||
if len(body) < 2:
|
||
return None
|
||
|
||
def get_pt(idx: int) -> tuple[float, float] | None:
|
||
if idx >= len(body):
|
||
return None
|
||
x, y, c = body[idx]
|
||
return (x, y) if c > 0 else None
|
||
|
||
# 1. Neck (index 1)
|
||
neck = get_pt(1)
|
||
if neck is not None:
|
||
return neck
|
||
|
||
# 2. Both shoulders (2, 5)
|
||
r_sh = get_pt(2)
|
||
l_sh = get_pt(5)
|
||
if r_sh is not None and l_sh is not None:
|
||
return ((r_sh[0] + l_sh[0]) / 2, (r_sh[1] + l_sh[1]) / 2)
|
||
|
||
# 3. Single shoulder
|
||
if r_sh is not None:
|
||
return r_sh
|
||
if l_sh is not None:
|
||
return l_sh
|
||
|
||
return None
|
||
|
||
|
||
def _position_keypoints_by_anchor(
|
||
rotated: dict[str, list[tuple[float, float, float]]],
|
||
original: dict[str, list[tuple[float, float, float]]],
|
||
) -> dict[str, list[tuple[float, float, float]]]:
|
||
"""
|
||
Translate rotated keypoints so the anchor (neck or shoulders) matches the original position.
|
||
Keeps the figure anchored in place based on its pre-rotation position.
|
||
"""
|
||
orig_anchor = _get_anchor_point(original)
|
||
rot_anchor = _get_anchor_point(rotated)
|
||
if orig_anchor is None or rot_anchor is None:
|
||
return rotated
|
||
|
||
dx = orig_anchor[0] - rot_anchor[0]
|
||
dy = orig_anchor[1] - rot_anchor[1]
|
||
|
||
result: dict[str, list[tuple[float, float, float]]] = {}
|
||
for part, pts in rotated.items():
|
||
result[part] = [(x + dx, y + dy, c) for x, y, c in pts]
|
||
return result
|
||
|
||
|
||
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 _draw_dashed_line(
|
||
canvas: np.ndarray,
|
||
p0: tuple[int, int],
|
||
p1: tuple[int, int],
|
||
color: tuple[int, int, int],
|
||
thickness: int = 1,
|
||
dash_len: int = 8,
|
||
gap_len: int = 6,
|
||
) -> None:
|
||
"""Draw a dashed line through p0 and p1, extended to canvas edges in both directions."""
|
||
h, w = canvas.shape[:2]
|
||
x0, y0 = float(p0[0]), float(p0[1])
|
||
x1, y1 = float(p1[0]), float(p1[1])
|
||
dx, dy = x1 - x0, y1 - y0
|
||
length = math.hypot(dx, dy)
|
||
if length < 1:
|
||
return
|
||
ux, uy = dx / length, dy / length
|
||
|
||
def t_to_edge(ox: float, oy: float, vx: float, vy: float) -> float:
|
||
t = float("inf")
|
||
if vx > 1e-9:
|
||
t = min(t, (w - 1 - ox) / vx)
|
||
elif vx < -1e-9:
|
||
t = min(t, -ox / vx)
|
||
if vy > 1e-9:
|
||
t = min(t, (h - 1 - oy) / vy)
|
||
elif vy < -1e-9:
|
||
t = min(t, -oy / vy)
|
||
return max(t, 0.0)
|
||
|
||
t_fwd = t_to_edge(x0, y0, ux, uy)
|
||
t_bwd = t_to_edge(x0, y0, -ux, -uy)
|
||
sx = x0 - t_bwd * ux
|
||
sy = y0 - t_bwd * uy
|
||
total = t_fwd + t_bwd
|
||
|
||
t = 0.0
|
||
drawing = True
|
||
while t < total:
|
||
seg = dash_len if drawing else gap_len
|
||
t_end = min(t + seg, total)
|
||
if drawing:
|
||
fx, fy = sx + t * ux, sy + t * uy
|
||
tx, ty = sx + t_end * ux, sy + t_end * uy
|
||
cv2.line(canvas, (int(fx), int(fy)), (int(tx), int(ty)), color, thickness)
|
||
t = t_end
|
||
drawing = not drawing
|
||
|
||
|
||
def render_openpose_image(
|
||
keypoints: dict[str, list[tuple[float, float, float]]],
|
||
width: int,
|
||
height: int,
|
||
debug: bool = False,
|
||
axis_line: tuple[tuple[int, int], tuple[int, int]] | None = None,
|
||
draw_face_mesh: bool = True,
|
||
) -> np.ndarray:
|
||
"""
|
||
Draw OpenPose skeleton on blank canvas. Black background (matches ControlNet pose images).
|
||
Uses static connection and draw-order definitions; no depth computation.
|
||
When debug=True and axis_line is provided, draws the rotation axis as a white dashed line
|
||
through the two given 2D points (extended to canvas edges).
|
||
draw_face_mesh: when False, the 70-point face mesh is omitted from the output image.
|
||
"""
|
||
canvas = np.zeros((height, width, 3), dtype=np.uint8)
|
||
|
||
body = keypoints.get("body", [])
|
||
|
||
# Build limb list from static definitions; filter by validity
|
||
max_limb_ratio = 0.5 # Max limb length vs image diagonal; filters bad connections
|
||
diag = (width**2 + height**2) ** 0.5
|
||
max_limb_len = diag * max_limb_ratio
|
||
|
||
limb_data: list[tuple[int, int, int, int, float]] = []
|
||
for i, ((a, b), layer, color_idx) in enumerate(BODY_LIMBS):
|
||
if a >= len(body) or b >= len(body):
|
||
continue
|
||
if a not in MAIN_BODY_INDICES or b not in MAIN_BODY_INDICES:
|
||
continue
|
||
x1, y1, c1 = body[a]
|
||
x2, y2, c2 = body[b]
|
||
if c1 <= 0 or c2 <= 0:
|
||
continue
|
||
limb_len = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5
|
||
if limb_len > max_limb_len:
|
||
continue
|
||
# Use max_x: limbs extending right (front when rotated) drawn last
|
||
max_x = max(x1, x2)
|
||
limb_data.append((a, b, layer, color_idx, max_x))
|
||
|
||
# Draw limbs: neutral gray (color is on nodes for debugging)
|
||
limb_data.sort(key=lambda t: (t[2], t[4]))
|
||
limb_color = (128, 128, 128) # gray
|
||
for a, b, _, _, _ in limb_data:
|
||
x1, y1, c1 = body[a]
|
||
x2, y2, c2 = body[b]
|
||
cv2.line(canvas, (int(x1), int(y1)), (int(x2), int(y2)), limb_color, 2)
|
||
|
||
# Draw joints: colored by index, twice as large for debugging
|
||
joint_radius = 8
|
||
joint_data = [
|
||
(i, x, y, c)
|
||
for i, (x, y, c) in enumerate(body)
|
||
if c > 0 and i in JOINT_LAYERS
|
||
]
|
||
joint_data.sort(key=lambda t: (JOINT_LAYERS.get(t[0], (0, 0))[0], t[1]))
|
||
for i, x, y, c in joint_data:
|
||
_, color_idx = JOINT_LAYERS.get(i, (0, 0))
|
||
color = BODY_COLORS[color_idx]
|
||
cv2.circle(canvas, (int(x), int(y)), joint_radius, color, -1)
|
||
|
||
# Debug: draw rotation axis (white, dashed, thin)
|
||
if debug and axis_line is not None:
|
||
_draw_dashed_line(canvas, axis_line[0], axis_line[1], (255, 255, 255), thickness=1)
|
||
|
||
# Face mesh (OpenPose 70-point) - draw before hands so body/face compose first
|
||
face_pts = keypoints.get("face", [])
|
||
if draw_face_mesh and len(face_pts) >= 2:
|
||
face_color = (180, 180, 180) # light gray
|
||
face_max_len = diag * 0.4
|
||
for a, b in FACE_EDGES:
|
||
if a >= len(face_pts) or b >= len(face_pts):
|
||
continue
|
||
x1, y1, c1 = face_pts[a]
|
||
x2, y2, c2 = face_pts[b]
|
||
if c1 <= 0 or c2 <= 0:
|
||
continue
|
||
limb_len = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5
|
||
if limb_len > face_max_len:
|
||
continue
|
||
cv2.line(canvas, (int(x1), int(y1)), (int(x2), int(y2)), face_color, 1)
|
||
# Draw face keypoints (small circles)
|
||
for x, y, c in face_pts:
|
||
if c > 0:
|
||
cv2.circle(canvas, (int(x), int(y)), 2, face_color, -1)
|
||
|
||
# Hands (21 keypoints each)
|
||
hand_color = (150, 150, 150)
|
||
hand_max_len = diag * 0.3
|
||
for hand_pts in [keypoints.get("hand_left", []), keypoints.get("hand_right", [])]:
|
||
if len(hand_pts) < 2:
|
||
continue
|
||
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
|
||
limb_len = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5
|
||
if limb_len > hand_max_len:
|
||
continue
|
||
cv2.line(canvas, (int(x1), int(y1)), (int(x2), int(y2)), hand_color, 1)
|
||
for x, y, c in hand_pts:
|
||
if c > 0:
|
||
cv2.circle(canvas, (int(x), int(y)), 2, hand_color, -1)
|
||
|
||
return canvas
|
||
|
||
|
||
def rotate_openpose(
|
||
image: np.ndarray,
|
||
pose_keypoint: Any | None,
|
||
direction: str,
|
||
degrees: float,
|
||
image_index: int = 0,
|
||
debug: bool = False,
|
||
mode: str = "simple",
|
||
recenter: bool = False,
|
||
draw_face_mesh: bool = True,
|
||
) -> tuple[np.ndarray, bool, dict[str, list[tuple[float, float, float]]] | None]:
|
||
"""
|
||
Main pipeline: extract/parse keypoints, detect torso, rotate, render.
|
||
Returns (output_image, success, rotated_keypoints). On failure, returns (input_image, False, None).
|
||
rotated_keypoints is internal format: {"body": [(x,y,c),...], "hand_left": [...], ...}
|
||
image_index: used for console logging when processing batches.
|
||
|
||
mode: "simple" uses orthographic projection at eye level (perspective=0).
|
||
"advanced" uses perspective projection with camera parameters derived from
|
||
ADVANCED_CAMERA_ELEVATION_DEG (35°) and ADVANCED_FIGURE_FILL (90%), which
|
||
gives geometrically correct depth inference and focal length for the assumed pose.
|
||
recenter: when True, horizontally centers the rotated figure within the image after all
|
||
other positioning corrections. Centering is measured from the widest body
|
||
keypoints; vertical position is not changed.
|
||
draw_face_mesh: when False, the 70-point face mesh is omitted from the rendered output.
|
||
"""
|
||
h, w = image.shape[:2]
|
||
|
||
# Derive camera parameters from mode and fixed assumptions
|
||
if mode == "advanced":
|
||
perspective, focal_length = _compute_advanced_params(h)
|
||
else:
|
||
perspective = 0.0
|
||
focal_length = 800.0 # unused in simple (orthographic) mode
|
||
|
||
# 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, None
|
||
|
||
# Scale normalized coords to image size (ComfyUI may pass 0-1 range)
|
||
keypoints = _scale_keypoints_to_image(keypoints, w, h)
|
||
|
||
# --- Stance detection (always logged) ---
|
||
# Compares the y-positions of both ankles relative to figure height.
|
||
# Within UPRIGHT_ANKLE_THRESHOLD → both feet on ground (spin axis = spine).
|
||
# Outside threshold → one foot on ground (spin axis = that ankle).
|
||
_body_kp = keypoints.get("body", [])
|
||
|
||
def _stance_pt(idx: int) -> tuple[float, float, float] | None:
|
||
if idx >= len(_body_kp):
|
||
return None
|
||
x, y, c = _body_kp[idx]
|
||
return (x, y, c) if c > 0 else None
|
||
|
||
_neck_s = _stance_pt(1)
|
||
_r_ank_s = _stance_pt(10) # COCO-18: r_ankle=10
|
||
_l_ank_s = _stance_pt(13) # COCO-18: l_ankle=13
|
||
|
||
stance: str = "unknown"
|
||
ground_ankle_pt: tuple[float, float, float] | None = None
|
||
ground_ankle_idx: int | None = None
|
||
|
||
if _neck_s is not None and _r_ank_s is not None and _l_ank_s is not None:
|
||
_lower_y = max(_r_ank_s[1], _l_ank_s[1])
|
||
_fig_h = abs(_lower_y - _neck_s[1])
|
||
if _fig_h >= 1:
|
||
_ank_diff = abs(_r_ank_s[1] - _l_ank_s[1])
|
||
if _ank_diff <= UPRIGHT_ANKLE_THRESHOLD * _fig_h:
|
||
stance = "both_feet"
|
||
print(f"[OpenPose Rotator] Image {image_index}: both feet on ground.")
|
||
else:
|
||
stance = "one_foot"
|
||
# Ground foot = the lower ankle (larger y, since y increases downward)
|
||
if _r_ank_s[1] >= _l_ank_s[1]:
|
||
ground_ankle_pt = _r_ank_s
|
||
ground_ankle_idx = 10
|
||
else:
|
||
ground_ankle_pt = _l_ank_s
|
||
ground_ankle_idx = 13
|
||
print(
|
||
f"[OpenPose Rotator] Image {image_index}: one foot on ground "
|
||
f"(ankle idx {ground_ankle_idx} at "
|
||
f"{ground_ankle_pt[0]:.0f}, {ground_ankle_pt[1]:.0f})."
|
||
)
|
||
|
||
# Console: pre-rotated figure (before rotation, after scaling) - only when debug
|
||
if debug:
|
||
pre_dict = _keypoints_to_openpose_dict(keypoints)
|
||
face_pts = keypoints.get("face", [])
|
||
face_visible = sum(1 for _, _, c in face_pts if c > 0)
|
||
if face_visible > 0:
|
||
print(f"[OpenPose Rotator] Image {image_index}: face mesh present ({face_visible}/{len(face_pts)} visible keypoints).")
|
||
else:
|
||
print(f"[OpenPose Rotator] Image {image_index}: no face mesh in input.")
|
||
print(f"[OpenPose Rotator] Image {image_index} PRE-ROTATED (direction={direction}, degrees={degrees}):")
|
||
print(f" body: {keypoints.get('body', [])}")
|
||
print(f" pose_keypoints_2d (flat): {pre_dict.get('people', [{}])[0].get('pose_keypoints_2d', [])}")
|
||
|
||
# Torso pivot — for one-foot stance, shift the rotation axis to the ground ankle's x
|
||
torso_pivot = compute_torso_pivot(keypoints)
|
||
if torso_pivot is None:
|
||
return image, False, None
|
||
|
||
if stance == "one_foot" and ground_ankle_pt is not None:
|
||
# cx = ground ankle; cy stays at torso level for depth-inference in z-calculation
|
||
pivot = (ground_ankle_pt[0], torso_pivot[1])
|
||
else:
|
||
pivot = torso_pivot
|
||
|
||
# Rotate keypoints around the chosen pivot
|
||
rotated = rotate_keypoints_3d(
|
||
keypoints,
|
||
pivot,
|
||
degrees,
|
||
direction,
|
||
mode=mode,
|
||
perspective=perspective,
|
||
focal_length=focal_length,
|
||
)
|
||
rotated = join_broken_segments(rotated, 0.0)
|
||
|
||
# Position by anchor
|
||
# • both_feet / unknown → align neck/shoulders (original behaviour)
|
||
# • one_foot → lock the ground ankle so it stays pinned to its original position
|
||
orig_anchor = _get_anchor_point(keypoints)
|
||
rot_anchor = _get_anchor_point(rotated)
|
||
|
||
if stance == "one_foot" and ground_ankle_pt is not None and ground_ankle_idx is not None:
|
||
body_rot = rotated.get("body", [])
|
||
if ground_ankle_idx < len(body_rot) and body_rot[ground_ankle_idx][2] > 0:
|
||
rot_ax, rot_ay, _ = body_rot[ground_ankle_idx]
|
||
dx = ground_ankle_pt[0] - rot_ax
|
||
dy = ground_ankle_pt[1] - rot_ay
|
||
rotated = {
|
||
part: [(px + dx, py + dy, c) for px, py, c in pts]
|
||
for part, pts in rotated.items()
|
||
}
|
||
else:
|
||
rotated = _position_keypoints_by_anchor(rotated, keypoints)
|
||
else:
|
||
rotated = _position_keypoints_by_anchor(rotated, keypoints)
|
||
|
||
# Horizontal in-frame correction for one-foot stance.
|
||
# After pinning the ankle, the swept side of the figure can overflow the image edge.
|
||
# When that happens, shift the *entire* figure (ankle included) just far enough to
|
||
# bring the overflowing side back to the image boundary.
|
||
# Only body keypoints are used to measure overflow — hand / face extremities are ignored.
|
||
if stance == "one_foot":
|
||
body_xs = [x for x, y, c in rotated.get("body", []) if c > 0]
|
||
if body_xs:
|
||
min_bx, max_bx = min(body_xs), max(body_xs)
|
||
if min_bx < 0 and max_bx > w:
|
||
shift_x = 0.0 # figure wider than frame — no good single shift
|
||
elif min_bx < 0:
|
||
shift_x = -min_bx
|
||
elif max_bx > w:
|
||
shift_x = w - max_bx
|
||
else:
|
||
shift_x = 0.0
|
||
if abs(shift_x) > 0.5:
|
||
print(
|
||
f"[OpenPose Rotator] Image {image_index}: figure shifted "
|
||
f"{shift_x:+.0f}px horizontally to stay in frame."
|
||
)
|
||
rotated = {
|
||
part: [(px + shift_x, py, c) for px, py, c in pts]
|
||
for part, pts in rotated.items()
|
||
}
|
||
|
||
# Optional horizontal recentering.
|
||
# Shifts the entire figure so the midpoint between its leftmost and rightmost visible
|
||
# body keypoints aligns with the horizontal centre of the image.
|
||
# Vertical positions are not touched. Applied after all other positioning corrections.
|
||
if recenter:
|
||
body_xs = [x for x, y, c in rotated.get("body", []) if c > 0]
|
||
if body_xs:
|
||
fig_center_x = (min(body_xs) + max(body_xs)) / 2.0
|
||
shift_x = (w / 2.0) - fig_center_x
|
||
if abs(shift_x) > 0.5:
|
||
rotated = {
|
||
part: [(px + shift_x, py, c) for px, py, c in pts]
|
||
for part, pts in rotated.items()
|
||
}
|
||
|
||
# Build debug axis line (after rotation + anchor positioning)
|
||
axis_line = None
|
||
if debug:
|
||
body_out = rotated.get("body", [])
|
||
|
||
if stance == "both_feet":
|
||
# Spine axis: neck → midpoint of both ankles
|
||
neck_out = body_out[1] if len(body_out) > 1 else None
|
||
r_ankle_out = body_out[10] if len(body_out) > 10 else None
|
||
l_ankle_out = body_out[13] if len(body_out) > 13 else None
|
||
if (neck_out is not None and neck_out[2] > 0
|
||
and r_ankle_out is not None and r_ankle_out[2] > 0
|
||
and l_ankle_out is not None and l_ankle_out[2] > 0):
|
||
ankle_mid_x = (r_ankle_out[0] + l_ankle_out[0]) / 2
|
||
ankle_mid_y = (r_ankle_out[1] + l_ankle_out[1]) / 2
|
||
axis_line = (
|
||
(int(round(neck_out[0])), int(round(neck_out[1]))),
|
||
(int(round(ankle_mid_x)), int(round(ankle_mid_y))),
|
||
)
|
||
print(
|
||
f"[OpenPose Rotator] Image {image_index}: both feet — spine axis "
|
||
f"from neck ({neck_out[0]:.0f},{neck_out[1]:.0f}) "
|
||
f"to ankle mid ({ankle_mid_x:.0f},{ankle_mid_y:.0f})."
|
||
)
|
||
|
||
elif stance == "one_foot" and ground_ankle_idx is not None:
|
||
# Vertical axis overlapping the ground ankle node
|
||
ank_out = body_out[ground_ankle_idx] if ground_ankle_idx < len(body_out) else None
|
||
if ank_out is not None and ank_out[2] > 0:
|
||
ax = int(round(ank_out[0]))
|
||
ay = int(round(ank_out[1]))
|
||
axis_line = ((ax, 0), (ax, h))
|
||
print(
|
||
f"[OpenPose Rotator] Image {image_index}: one foot — vertical axis "
|
||
f"at ankle ({ax}, {ay})."
|
||
)
|
||
|
||
if axis_line is None and orig_anchor is not None and rot_anchor is not None:
|
||
# Fallback (stance unknown or required keypoints missing)
|
||
dx_fb = orig_anchor[0] - rot_anchor[0]
|
||
axis_xi = int(round(pivot[0] + dx_fb))
|
||
axis_line = ((axis_xi, 0), (axis_xi, h))
|
||
|
||
# Console: post-rotated figure (after rotation and anchor positioning) - only when debug
|
||
if debug:
|
||
post_dict = _keypoints_to_openpose_dict(rotated)
|
||
print(f"[OpenPose Rotator] Image {image_index} POST-ROTATED (direction={direction}, degrees={degrees}):")
|
||
print(f" body: {rotated.get('body', [])}")
|
||
print(f" pose_keypoints_2d (flat): {post_dict.get('people', [{}])[0].get('pose_keypoints_2d', [])}")
|
||
|
||
# Render using static connection and draw-order definitions
|
||
rendered = render_openpose_image(rotated, w, h, debug=debug, axis_line=axis_line, draw_face_mesh=draw_face_mesh)
|
||
return rendered, True, rotated
|