draw_bodypose
This commit is contained in:
parent
6ebb9e3aaf
commit
96ec70e47d
|
@ -0,0 +1,154 @@
|
|||
import json
|
||||
import numpy as np
|
||||
from typing import List
|
||||
import math
|
||||
import cv2
|
||||
|
||||
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})"
|
||||
|
||||
def is_normalized(keypoints: List[Keypoint]) -> bool:
|
||||
"""
|
||||
Check if the keypoints are normalized between 0 and 1.
|
||||
|
||||
Args:
|
||||
keypoints (List[Keypoint]): A list of Keypoint objects.
|
||||
|
||||
Returns:
|
||||
bool: True if all keypoints are normalized, False otherwise.
|
||||
"""
|
||||
for keypoint in keypoints:
|
||||
if not (0 <= keypoint.x <= 1 and 0 <= keypoint.y <= 1):
|
||||
return False
|
||||
return True
|
||||
|
||||
def draw_bodypose(canvas: np.ndarray, keypoints: List[Keypoint], 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 = 4
|
||||
|
||||
# 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
|
||||
|
||||
limbSeq = [
|
||||
[2, 3], [2, 6], [3, 4], [4, 5],
|
||||
[6, 7], [7, 8], [2, 9], [9, 10],
|
||||
[10, 11], [2, 12], [12, 13], [13, 14],
|
||||
[2, 1], [1, 15], [15, 17], [1, 16],
|
||||
[16, 18],
|
||||
]
|
||||
|
||||
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]]
|
||||
|
||||
for (k1_index, k2_index), color in zip(limbSeq, colors):
|
||||
keypoint1 = keypoints[k1_index - 1]
|
||||
keypoint2 = keypoints[k2_index - 1]
|
||||
|
||||
if keypoint1 is None or keypoint2 is None:
|
||||
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:
|
||||
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
|
||||
|
||||
def openpose_json_to_keypoints(json_file: str) -> List[Keypoint]:
|
||||
"""
|
||||
Convert OpenPose JSON file to a list of Keypoint objects.
|
||||
|
||||
Args:
|
||||
json_file (str): The path to the OpenPose JSON file.
|
||||
|
||||
Returns:
|
||||
List[Keypoint]: A list of Keypoint objects representing the keypoints extracted from the JSON file.
|
||||
"""
|
||||
with open(json_file, 'r') as file:
|
||||
data = json.load(file)
|
||||
keypoints = data[0]['people'][0]['pose_keypoints_2d']
|
||||
keypoints = [Keypoint(keypoints[i], keypoints[i + 1]) for i in range(0, len(keypoints), 3)]
|
||||
return keypoints
|
||||
|
||||
def main():
|
||||
# Create a blank canvas
|
||||
canvas = np.zeros((768, 768, 3), dtype=np.uint8)
|
||||
|
||||
# Load the keypoints from the OpenPose JSON file
|
||||
keypoints = openpose_json_to_keypoints('openpose_output.json')
|
||||
|
||||
# Draw the body pose on the canvas
|
||||
canvas = draw_bodypose(canvas, keypoints)
|
||||
|
||||
# Display the result
|
||||
cv2.imshow('Body Pose', canvas)
|
||||
cv2.waitKey(0)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
# Here is an example of input json file, create a List of Keypoint objects from the json file and draw the body pose on a canvas using the draw_bodypose function.
|
||||
#
|
||||
# [
|
||||
# {
|
||||
# "people": [
|
||||
# {
|
||||
# "pose_keypoints_2d": [423.575046753433, 113.537678172812, 1, 378.126093103822, 185.044031914867, 1, 416.303214169495, 169.288394649668, 1, 470.841958549029, 122.021482854073, 1, 519.320842441947, 48.0911849173717, 1, 339.948972038149, 200.799669180065, 1, 322.981362675627, 289.273632284642, 1, 316.921502189012, 391.079288459771, 1, 439.330684018632, 363.203930221343, 1, 412.667297877527, 541.363828527819, 1, 375.096162860515, 702.556117471773, 1, 375.096162860515, 364.415902318666, 1, 359.340525595316, 530.456079651912, 1, 292.682060242553, 674.680759233345, 1, 411.455325780204, 105.053873491551, 1, 418.727158364141, 102.629929296905, 1, 383.579967541775, 120.80951075675, 1, 371.460246568546, 122.021482854073, 1]
|
||||
# }
|
||||
# ],
|
||||
# "canvas_height": 768,
|
||||
# "canvas_width": 768
|
||||
# }
|
||||
# ]
|
|
@ -0,0 +1,11 @@
|
|||
[
|
||||
{
|
||||
"people": [
|
||||
{
|
||||
"pose_keypoints_2d": [423.575046753433, 113.537678172812, 1, 378.126093103822, 185.044031914867, 1, 416.303214169495, 169.288394649668, 1, 470.841958549029, 122.021482854073, 1, 519.320842441947, 48.0911849173717, 1, 339.948972038149, 200.799669180065, 1, 322.981362675627, 289.273632284642, 1, 316.921502189012, 391.079288459771, 1, 439.330684018632, 363.203930221343, 1, 412.667297877527, 541.363828527819, 1, 375.096162860515, 702.556117471773, 1, 375.096162860515, 364.415902318666, 1, 359.340525595316, 530.456079651912, 1, 292.682060242553, 674.680759233345, 1, 411.455325780204, 105.053873491551, 1, 418.727158364141, 102.629929296905, 1, 383.579967541775, 120.80951075675, 1, 371.460246568546, 122.021482854073, 1]
|
||||
}
|
||||
],
|
||||
"canvas_height": 768,
|
||||
"canvas_width": 768
|
||||
}
|
||||
]
|
Loading…
Reference in New Issue