186 lines
6.9 KiB
Python
186 lines
6.9 KiB
Python
from typing import List
|
|
import numpy as np
|
|
import math
|
|
import cv2
|
|
from numpy import ndarray
|
|
|
|
coco_limbSeq = [
|
|
[1, 2], [1, 5], [2, 3], [3, 4],
|
|
[5, 6], [6, 7], [1, 8], [8, 9],
|
|
[9, 10], [1, 11], [11, 12], [12, 13],
|
|
[1, 0], [0, 14], [14, 16], [0, 15],
|
|
[15, 17],
|
|
]
|
|
|
|
coco_colors = [
|
|
[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
|
|
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
|
|
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]
|
|
]
|
|
|
|
body_25_limbSeq = [
|
|
[1, 8], [1, 2], [1, 5], [2, 3],
|
|
[3, 4], [5, 6], [6, 7], [8, 9],
|
|
[9, 10], [10, 11], [8, 12], [12, 13],
|
|
[13, 14], [1, 0], [0, 15], [15, 17],
|
|
[0, 16], [16, 18], [14, 19], [19, 20],
|
|
[14, 21], [11, 22], [22, 23], [11, 24]
|
|
]
|
|
body_25_colors = [
|
|
[255, 0, 85], [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0],
|
|
[0, 255, 0], [255, 0, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255],
|
|
[0, 0, 255], [255, 0, 170], [170, 0, 255], [255, 0, 255], [85, 0, 255], [0, 0, 255], [0, 0, 255],
|
|
[0, 0, 255], [0, 0, 255], [0, 255, 255], [0, 255, 255], [0, 255, 255]
|
|
]
|
|
|
|
body_25B_limbSeq = [
|
|
[0, 1], [0, 2], [1, 3], [2, 4], [5, 7], [6, 8], [7, 9], [8, 10],
|
|
[5, 11], [6, 12], [11, 13], [12, 14], [13, 15], [14, 16], [15, 19],
|
|
[19, 20], [15, 21], [16, 22], [22, 23], [16, 24], [5, 17], [6, 17],
|
|
[17, 18], [11, 12]
|
|
]
|
|
|
|
body_25B_colors = [
|
|
[255, 0, 85], [170, 0, 255], [255, 0, 170], [85, 0, 255], [255, 0, 255],
|
|
[170, 255, 0], [255, 85, 0], [85, 255, 0], [255, 170, 0], [0, 255, 0],
|
|
[255, 255, 0], [0, 170, 255], [0, 255, 85], [0, 85, 255], [0, 255, 170],
|
|
[0, 0, 255], [0, 255, 255], [255, 0, 0], [255, 0, 0], [0, 0, 255],
|
|
[0, 0, 255], [0, 0, 255], [0, 255, 255], [0, 255, 255], [0, 255, 255]
|
|
]
|
|
|
|
class Keypoint:
|
|
def __init__(self, x: float, y: float, confidence: float = 1.0):
|
|
"""
|
|
Initialize a Keypoint object.
|
|
|
|
Args:
|
|
x (float): The x-coordinate of the keypoint.
|
|
y (float): The y-coordinate of the keypoint.
|
|
confidence (float): The confidence score of the keypoint. Default is 1.0.
|
|
"""
|
|
self.x = x
|
|
self.y = y
|
|
self.confidence = confidence
|
|
|
|
def __repr__(self):
|
|
return f"Keypoint(x={self.x}, y={self.y}, confidence={self.confidence})"
|
|
|
|
class Skeleton:
|
|
def __init__(self, keypoints: List[Keypoint]):
|
|
self.keypoints = keypoints
|
|
|
|
def __repr__(self):
|
|
return f"Skeleton(keypoints={self.keypoints})"
|
|
|
|
def is_healthy_skeleton(self):
|
|
for keypoint in self.keypoints:
|
|
if keypoint.confidence == 0.0:
|
|
return False
|
|
return True
|
|
|
|
class Skeleton_Seqence:
|
|
def __init__(self, skeletons: List[Skeleton]):
|
|
self.skeletons_frame = skeletons
|
|
|
|
def __repr__(self):
|
|
return f"Skeleton_Seqence(Skeleton_frames={self.skeletons_frame})"
|
|
|
|
def get_frame(self, frame_index: int) -> Skeleton:
|
|
return self.skeletons_frame[frame_index]
|
|
|
|
def add_frame(self, skeleton: Skeleton):
|
|
self.skeletons_frame.append(skeleton)
|
|
|
|
def is_healthy_seqence(self):
|
|
for skeleton in self.skeletons_frame:
|
|
if not skeleton.is_healthy_skeleton():
|
|
return False
|
|
return True
|
|
|
|
def fix_keypoints(keypoints, last_keypoints, next_keypoints):
|
|
if not keypoints or not last_keypoints or not next_keypoints:
|
|
return keypoints
|
|
for i, keypoint in enumerate(keypoints):
|
|
if keypoint.confidence == 0.0 and last_keypoints and next_keypoints:
|
|
last_keypoint = last_keypoints[i]
|
|
next_keypoint = next_keypoints[i]
|
|
if last_keypoint.confidence > 0 and next_keypoint.confidence > 0:
|
|
keypoint.x = (last_keypoint.x + next_keypoint.x) / 2
|
|
keypoint.y = (last_keypoint.y + next_keypoint.y) / 2
|
|
keypoint.confidence = (last_keypoint.confidence + next_keypoint.confidence) / 2
|
|
return keypoints
|
|
|
|
def get_time_slice_for_Skeleton_Seqences(skeleton_seqences: List[Skeleton_Seqence], frame_index: int) -> List[Skeleton]:
|
|
return [skeleton_seq.get_frame(frame_index) for skeleton_seq in skeleton_seqences]
|
|
|
|
|
|
def is_normalized(keypoints: List[Keypoint]) -> bool:
|
|
for keypoint in keypoints:
|
|
if not (0 <= keypoint.x <= 1 and 0 <= keypoint.y <= 1):
|
|
return False
|
|
return True
|
|
|
|
|
|
def draw_bodypose(canvas: ndarray, keypoints: List[Keypoint], limbSeq, colors, xinsr_stick_scaling: bool = False) -> np.ndarray:
|
|
"""
|
|
Draw keypoints and limbs representing body pose on a given canvas.
|
|
|
|
Args:
|
|
canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the body pose.
|
|
keypoints (List[Keypoint]): A list of Keypoint objects representing the body keypoints to be drawn.
|
|
xinsr_stick_scaling (bool): Whether or not scaling stick width for xinsr ControlNet
|
|
|
|
Returns:
|
|
np.ndarray: A 3D numpy array representing the modified canvas with the drawn body pose.
|
|
|
|
Note:
|
|
The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
|
|
"""
|
|
if not is_normalized(keypoints):
|
|
H, W = 1.0, 1.0
|
|
else:
|
|
H, W, _ = canvas.shape
|
|
|
|
CH, CW, _ = canvas.shape
|
|
stickwidth = 2
|
|
|
|
# Ref: https://huggingface.co/xinsir/controlnet-openpose-sdxl-1.0
|
|
max_side = max(CW, CH)
|
|
if xinsr_stick_scaling:
|
|
stick_scale = 1 if max_side < 500 else min(2 + (max_side // 1000), 7)
|
|
else :
|
|
stick_scale = 1
|
|
|
|
if keypoints is None or len(keypoints) == 0:
|
|
return canvas
|
|
|
|
for (k1_index, k2_index), color in zip(limbSeq, colors):
|
|
keypoint1 = keypoints[k1_index]
|
|
keypoint2 = keypoints[k2_index]
|
|
|
|
if keypoint1 is None or keypoint2 is None or keypoint1.confidence == 0 or keypoint2.confidence == 0:
|
|
# if keypoint1 is None or keypoint1.confidence == 0:
|
|
# print(f"keypoint failed: {k1_index}")
|
|
# if keypoint2 is None or keypoint2.confidence == 0:
|
|
# print(f"keypoint failed: {k2_index}")
|
|
continue
|
|
|
|
Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
|
|
X = np.array([keypoint1.y, keypoint2.y]) * float(H)
|
|
mX = np.mean(X)
|
|
mY = np.mean(Y)
|
|
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
|
|
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
|
|
polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth*stick_scale), int(angle), 0, 360, 1)
|
|
cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
|
|
|
|
for keypoint, color in zip(keypoints, colors):
|
|
if keypoint is None or keypoint.confidence == 0:
|
|
continue
|
|
|
|
x, y = keypoint.x, keypoint.y
|
|
x = int(x * W)
|
|
y = int(y * H)
|
|
cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
|
|
|
|
return canvas |