Skip to content

Commit

Permalink
Merge pull request #381 from huchenlei/render
Browse files Browse the repository at this point in the history
Add node to render openpose kps JSON
Fannovel16 authored Jun 21, 2024
2 parents 36fca60 + 6c95d1c commit 4b18c70
Showing 3 changed files with 161 additions and 29 deletions.
72 changes: 69 additions & 3 deletions node_wrappers/pose_keypoint_postprocess.py
Original file line number Diff line number Diff line change
@@ -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)",
}
117 changes: 91 additions & 26 deletions src/controlnet_aux/dwpose/__init__.py
Original file line number Diff line number Diff line change
@@ -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,
)
1 change: 1 addition & 0 deletions src/controlnet_aux/dwpose/types.py
Original file line number Diff line number Diff line change
@@ -20,6 +20,7 @@ class BodyResult(NamedTuple):

HandResult = List[Keypoint]
FaceResult = List[Keypoint]
AnimalPoseResult = List[Keypoint]


class PoseResult(NamedTuple):

0 comments on commit 4b18c70

Please sign in to comment.