From 96ec70e47df6f6bb9c72983d8d2b1d31abf44c68 Mon Sep 17 00:00:00 2001 From: zaqxs123456 <18200842@life.hkbu.edu.hk> Date: Mon, 23 Sep 2024 18:03:14 +0800 Subject: [PATCH] draw_bodypose --- openpose_gen.py | 154 +++++++++++++++++++++++++++++++++++++++++++ openpose_output.json | 11 ++++ 2 files changed, 165 insertions(+) create mode 100644 openpose_gen.py create mode 100644 openpose_output.json diff --git a/openpose_gen.py b/openpose_gen.py new file mode 100644 index 0000000..eb21585 --- /dev/null +++ b/openpose_gen.py @@ -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 +# } +# ] \ No newline at end of file diff --git a/openpose_output.json b/openpose_output.json new file mode 100644 index 0000000..a3a74f1 --- /dev/null +++ b/openpose_output.json @@ -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 + } +] \ No newline at end of file