draw_bodypose

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
+#   }
+# ]

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
+  }
+]