Add node to render openpose kps JSON

node_wrappers/pose_keypoint_postprocess.py

@@ -8,6 +8,9 @@
 import torch
 import itertools
 
+from ..src.controlnet_aux.dwpose import draw_poses, draw_animalposes, decode_json_as_poses
+
+
 """
 Format of POSE_KEYPOINT (AP10K keypoints):
 [{
@@ -252,23 +255,86 @@ def convert(self, pose_kps, id_include, **parts_width_height):
                 part_bboxes = processor.get_xyxy_bboxes(part_name, bbox_size)
                 id_coordinates = {idx: part_bbox+(processor.width, processor.height) for idx, part_bbox in part_bboxes.items()}
                 tracked[part_name][person_idx] = id_coordinates
-        
+
         for class_name, class_data in tracked.items():
             for class_id in class_data.keys():
                 class_id_str = str(class_id)
                 # Use the incoming prompt for each class name and ID
                 _class_name = class_name.replace('L', '').replace('R', '').lower()
                 prompt_string += f'"{class_id_str}.{class_name}": "({_class_name})",\n'
-        
+
         return (tracked, prompt_string)
 
+
+def numpy2torch(np_image: np.ndarray) -> torch.Tensor:
+    """ [H, W, C] => [B=1, H, W, C]"""
+    return torch.from_numpy(np_image.astype(np.float32) / 255).unsqueeze(0)
+
+
+class RenderPeopleKps:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "kps": ("POSE_KEYPOINT",),
+                "render_body": ("BOOLEAN", {"default": True}),
+                "render_hand": ("BOOLEAN", {"default": True}),
+                "render_face": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "render"
+    CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
+
+    def render(self, kps, render_body, render_hand, render_face) -> tuple[np.ndarray]:
+        if isinstance(kps, list):
+            kps = kps[0]
+
+        poses, _, height, width = decode_json_as_poses(kps)
+        np_image = draw_poses(
+            poses,
+            height,
+            width,
+            render_body,
+            render_hand,
+            render_face,
+        )
+        return (numpy2torch(np_image),)
+
+class RenderAnimalKps:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "kps": ("POSE_KEYPOINT",),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "render"
+    CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
+
+    def render(self, kps) -> tuple[np.ndarray]:
+        if isinstance(kps, list):
+            kps = kps[0]
+
+        _, poses, height, width = decode_json_as_poses(kps)
+        np_image = draw_animalposes(poses, height, width)
+        return (numpy2torch(np_image),)
+
+
 NODE_CLASS_MAPPINGS = {
     "SavePoseKpsAsJsonFile": SavePoseKpsAsJsonFile,
     "FacialPartColoringFromPoseKps": FacialPartColoringFromPoseKps,
-    "UpperBodyTrackingFromPoseKps": UpperBodyTrackingFromPoseKps
+    "UpperBodyTrackingFromPoseKps": UpperBodyTrackingFromPoseKps,
+    "RenderPeopleKps": RenderPeopleKps,
+    "RenderAnimalKps": RenderAnimalKps,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "SavePoseKpsAsJsonFile": "Save Pose Keypoints",
     "FacialPartColoringFromPoseKps": "Colorize Facial Parts from PoseKPS",
     "UpperBodyTrackingFromPoseKps": "Upper Body Tracking From PoseKps (InstanceDiffusion)",
+    "RenderPeopleKps": "Render Pose JSON (Human)",
+    "RenderAnimalKps": "Render Pose JSON (Animal)",
 }

src/controlnet_aux/dwpose/__init__.py

@@ -16,7 +16,7 @@
 from .body import Body, BodyResult, Keypoint
 from .hand import Hand
 from .face import Face
-from .types import PoseResult, HandResult, FaceResult
+from .types import PoseResult, HandResult, FaceResult, AnimalPoseResult
 from huggingface_hub import hf_hub_download
 from .wholebody import Wholebody
 import warnings
@@ -27,6 +27,70 @@
 
 from typing import Tuple, List, Callable, Union, Optional
 
+
+def draw_animalposes(animals: list[list[Keypoint]], H: int, W: int) -> np.ndarray:
+    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+    for animal_pose in animals:
+        canvas = draw_animalpose(canvas, animal_pose)
+    return canvas
+
+
+def draw_animalpose(canvas: np.ndarray, keypoints: list[Keypoint]) -> np.ndarray:
+    # order of the keypoints for AP10k and a standardized list of colors for limbs
+    keypointPairsList = [
+        (1, 2),
+        (2, 3),
+        (1, 3),
+        (3, 4),
+        (4, 9),
+        (9, 10),
+        (10, 11),
+        (4, 6),
+        (6, 7),
+        (7, 8),
+        (4, 5),
+        (5, 15),
+        (15, 16),
+        (16, 17),
+        (5, 12),
+        (12, 13),
+        (13, 14),
+    ]
+    colorsList = [
+        (255, 255, 255),
+        (100, 255, 100),
+        (150, 255, 255),
+        (100, 50, 255),
+        (50, 150, 200),
+        (0, 255, 255),
+        (0, 150, 0),
+        (0, 0, 255),
+        (0, 0, 150),
+        (255, 50, 255),
+        (255, 0, 255),
+        (255, 0, 0),
+        (150, 0, 0),
+        (255, 255, 100),
+        (0, 150, 0),
+        (255, 255, 0),
+        (150, 150, 150),
+    ]  # 16 colors needed
+
+    for ind, (i, j) in enumerate(keypointPairsList):
+        p1 = keypoints[i - 1]
+        p2 = keypoints[j - 1]
+
+        if p1 is not None and p2 is not None:
+            cv2.line(
+                canvas,
+                (int(p1.x), int(p1.y)),
+                (int(p2.x), int(p2.y)),
+                colorsList[ind],
+                5,
+            )
+    return canvas
+
+
 def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, draw_face=True):
     """
     Draw the detected poses on an empty canvas.
@@ -58,35 +122,36 @@ def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, dr
     return canvas
 
 
-def decode_json_as_poses(json_string: str, normalize_coords: bool = False) -> Tuple[List[PoseResult], int, int]:
-    """ Decode the json_string complying with the openpose JSON output format
+def decode_json_as_poses(
+    pose_json: dict,
+) -> Tuple[List[PoseResult], List[AnimalPoseResult], int, int]:
+    """Decode the json_string complying with the openpose JSON output format
     to poses that controlnet recognizes.
     https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md
 
     Args:
         json_string: The json string to decode.
-        normalize_coords: Whether to normalize coordinates of each keypoint by canvas height/width.
-                          `draw_pose` only accepts normalized keypoints. Set this param to True if
-                          the input coords are not normalized.
-
+
     Returns:
-        poses
+        human_poses
+        animal_poses
         canvas_height
-        canvas_width                      
+        canvas_width
     """
-    pose_json = json.loads(json_string)
-    height = pose_json['canvas_height']
-    width = pose_json['canvas_width']
+    height = pose_json["canvas_height"]
+    width = pose_json["canvas_width"]
 
     def chunks(lst, n):
         """Yield successive n-sized chunks from lst."""
         for i in range(0, len(lst), n):
-            yield lst[i:i + n]
-
-    def decompress_keypoints(numbers: Optional[List[float]]) -> Optional[List[Optional[Keypoint]]]:
+            yield lst[i : i + n]
+
+    def decompress_keypoints(
+        numbers: Optional[List[float]],
+    ) -> Optional[List[Optional[Keypoint]]]:
         if not numbers:
             return None
-        
+
         assert len(numbers) % 3 == 0
 
         def create_keypoint(x, y, c):
@@ -95,21 +160,21 @@ def create_keypoint(x, y, c):
             keypoint = Keypoint(x, y)
             return keypoint
 
-        return [
-            create_keypoint(x, y, c)
-            for x, y, c in chunks(numbers, n=3)
-        ]
-
+        return [create_keypoint(x, y, c) for x, y, c in chunks(numbers, n=3)]
+
     return (
         [
             PoseResult(
-                body=BodyResult(keypoints=decompress_keypoints(pose.get('pose_keypoints_2d'))),
-                left_hand=decompress_keypoints(pose.get('hand_left_keypoints_2d')),
-                right_hand=decompress_keypoints(pose.get('hand_right_keypoints_2d')),
-                face=decompress_keypoints(pose.get('face_keypoints_2d'))
+                body=BodyResult(
+                    keypoints=decompress_keypoints(pose.get("pose_keypoints_2d"))
+                ),
+                left_hand=decompress_keypoints(pose.get("hand_left_keypoints_2d")),
+                right_hand=decompress_keypoints(pose.get("hand_right_keypoints_2d")),
+                face=decompress_keypoints(pose.get("face_keypoints_2d")),
             )
-            for pose in pose_json['people']
+            for pose in pose_json.get("people", [])
         ],
+        [decompress_keypoints(pose) for pose in pose_json.get("animals", [])],
         height,
         width,
     )

src/controlnet_aux/dwpose/types.py

@@ -20,6 +20,7 @@ class BodyResult(NamedTuple):
 
 HandResult = List[Keypoint]
 FaceResult = List[Keypoint]
+AnimalPoseResult = List[Keypoint]
 
 
 class PoseResult(NamedTuple):