diff --git a/pose_utils.py b/pose_utils.py index 90d3fc7..cd54f4c 100644 --- a/pose_utils.py +++ b/pose_utils.py @@ -445,70 +445,61 @@ def rotate_keypoints_3d( return result -def _fit_and_center_keypoints_on_canvas( - keypoints: dict[str, list[tuple[float, float, float]]], - width: int, - height: int, - padding: float = 0.05, +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]]]: """ - Scale keypoints to fit within canvas (if they exceed bounds), then center. - Prevents feet/head from being cut off when the figure is taller than the image. - Uses main body only (excludes face indices) for centroid/bbox so face points - don't pull the figure down when centered. + 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. """ - all_x, all_y, total_conf = 0.0, 0.0, 0.0 - xs, ys = [], [] - for part, pts in keypoints.items(): - for i, (x, y, c) in enumerate(pts): - if c <= 0: - continue - if part == "body" and i in FACE_INDICES: - continue # Exclude face from centroid/bbox - center on main body - all_x += x * c - all_y += y * c - total_conf += c - xs.append(x) - ys.append(y) + orig_anchor = _get_anchor_point(original) + rot_anchor = _get_anchor_point(rotated) + if orig_anchor is None or rot_anchor is None: + return rotated - if total_conf <= 0 or not xs or not ys: - return keypoints - - min_x, max_x = min(xs), max(xs) - min_y, max_y = min(ys), max(ys) - - # Centroid (stays fixed when scaling around it) - cx = all_x / total_conf - cy = all_y / total_conf - - # Max extent from centroid in each direction (we center the centroid, so we must fit - # the larger of the two sides; using bbox_h/2 fails when centroid is off-center) - extent_x = max(cx - min_x, max_x - cx) - extent_y = max(cy - min_y, max_y - cy) - - # Usable canvas with padding (half-extent each side of center) - pad_w = width * padding - pad_h = height * padding - usable_half_w = (width - 2 * pad_w) / 2.0 - usable_half_h = (height - 2 * pad_h) / 2.0 - - # Scale down so extent from centroid fits within usable half-canvas - scale = 1.0 - if extent_x > usable_half_w and usable_half_w > 0: - scale = min(scale, usable_half_w / extent_x) - if extent_y > usable_half_h and usable_half_h > 0: - scale = min(scale, usable_half_h / extent_y) - target_x = width / 2.0 - target_y = height / 2.0 - dx = target_x - cx - dy = target_y - cy + 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 keypoints.items(): - result[part] = [ - (cx + (x - cx) * scale + dx, cy + (y - cy) * scale + dy, c) - for x, y, c in pts - ] + for part, pts in rotated.items(): + result[part] = [(x + dx, y + dy, c) for x, y, c in pts] return result @@ -633,10 +624,10 @@ def rotate_openpose( rotated = rotate_keypoints_3d(keypoints, pivot, degrees, direction) rotated = join_broken_segments(rotated, 0.0) - # Scale to fit and center figure (prevents feet/head cutoff, fixes right/left bias) - rotated = _fit_and_center_keypoints_on_canvas(rotated, w, h) + # Position by anchor: align neck/shoulders with original figure position + rotated = _position_keypoints_by_anchor(rotated, keypoints) - # Console: post-rotated figure (after rotation, fit, and center) + # Console: post-rotated figure (after rotation and anchor positioning) 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', [])}")