import folder_paths import json import os import numpy as np import cv2 from PIL import ImageColor from einops import rearrange import torch import itertools from ..src.custom_controlnet_aux.dwpose import draw_poses, draw_animalposes, decode_json_as_poses """ Format of POSE_KEYPOINT (AP10K keypoints): [{ "version": "ap10k", "animals": [ [[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]], [[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]], ... ], "canvas_height": 512, "canvas_width": 768 },...] Format of POSE_KEYPOINT (OpenPose keypoints): [{ "people": [ { 'pose_keypoints_2d': [[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]] "face_keypoints_2d": [[x1, y1, 1], [x2, y2, 1],..., [x68, y68, 1]], "hand_left_keypoints_2d": [[x1, y1, 1], [x2, y2, 1],..., [x21, y21, 1]], "hand_right_keypoints_2d":[[x1, y1, 1], [x2, y2, 1],..., [x21, y21, 1]], } ], "canvas_height": canvas_height, "canvas_width": canvas_width, },...] """ class SavePoseKpsAsJsonFile: @classmethod def INPUT_TYPES(s): return { "required": { "pose_kps": ("POSE_KEYPOINT",), "filename_prefix": ("STRING", {"default": "PoseKeypoint"}) } } RETURN_TYPES = () FUNCTION = "save_pose_kps" OUTPUT_NODE = True CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess" def __init__(self): self.output_dir = folder_paths.get_output_directory() self.type = "output" self.prefix_append = "" def save_pose_kps(self, pose_kps, filename_prefix): filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = \ folder_paths.get_save_image_path(filename_prefix, self.output_dir, pose_kps[0]["canvas_width"], pose_kps[0]["canvas_height"]) file = f"{filename}_{counter:05}.json" with open(os.path.join(full_output_folder, file), 'w') as f: json.dump(pose_kps , f) return {} #COCO-Wholebody doesn't have eyebrows as it inherits 68 keypoints format #Perhaps eyebrows can be estimated tho FACIAL_PARTS = ["skin", "left_eye", "right_eye", "nose", "upper_lip", "inner_mouth", "lower_lip"] LAPA_COLORS = dict( skin="rgb(0, 153, 255)", left_eye="rgb(0, 204, 153)", right_eye="rgb(255, 153, 0)", nose="rgb(255, 102, 255)", upper_lip="rgb(102, 0, 51)", inner_mouth="rgb(255, 204, 255)", lower_lip="rgb(255, 0, 102)" ) #One-based index def kps_idxs(start, end): step = -1 if start > end else 1 return list(range(start-1, end+1-1, step)) #Source: https://www.researchgate.net/profile/Fabrizio-Falchi/publication/338048224/figure/fig1/AS:837860722741255@1576772971540/68-facial-landmarks.jpg FACIAL_PART_RANGES = dict( skin=kps_idxs(1, 17) + kps_idxs(27, 18), nose=kps_idxs(28, 36), left_eye=kps_idxs(37, 42), right_eye=kps_idxs(43, 48), upper_lip=kps_idxs(49, 55) + kps_idxs(65, 61), lower_lip=kps_idxs(61, 68), inner_mouth=kps_idxs(61, 65) + kps_idxs(55, 49) ) def is_normalized(keypoints) -> bool: point_normalized = [ 0 <= np.abs(k[0]) <= 1 and 0 <= np.abs(k[1]) <= 1 for k in keypoints if k is not None ] if not point_normalized: return False return np.all(point_normalized) class FacialPartColoringFromPoseKps: @classmethod def INPUT_TYPES(s): input_types = { "required": {"pose_kps": ("POSE_KEYPOINT",), "mode": (["point", "polygon"], {"default": "polygon"})} } for facial_part in FACIAL_PARTS: input_types["required"][facial_part] = ("STRING", {"default": LAPA_COLORS[facial_part], "multiline": False}) return input_types RETURN_TYPES = ("IMAGE",) FUNCTION = "colorize" CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess" def colorize(self, pose_kps, mode, **facial_part_colors): pose_frames = pose_kps np_frames = [self.draw_kps(pose_frame, mode, **facial_part_colors) for pose_frame in pose_frames] np_frames = np.stack(np_frames, axis=0) return (torch.from_numpy(np_frames).float() / 255.,) def draw_kps(self, pose_frame, mode, **facial_part_colors): width, height = pose_frame["canvas_width"], pose_frame["canvas_height"] canvas = np.zeros((height, width, 3), dtype=np.uint8) for person, part_name in itertools.product(pose_frame["people"], FACIAL_PARTS): n = len(person["face_keypoints_2d"]) // 3 facial_kps = rearrange(np.array(person["face_keypoints_2d"]), "(n c) -> n c", n=n, c=3)[:, :2] if is_normalized(facial_kps): facial_kps *= (width, height) facial_kps = facial_kps.astype(np.int32) part_color = ImageColor.getrgb(facial_part_colors[part_name])[:3] part_contours = facial_kps[FACIAL_PART_RANGES[part_name], :] if mode == "point": for pt in part_contours: cv2.circle(canvas, pt, radius=2, color=part_color, thickness=-1) else: cv2.fillPoly(canvas, pts=[part_contours], color=part_color) return canvas # https://raw.githubusercontent.com/CMU-Perceptual-Computing-Lab/openpose/master/.github/media/keypoints_pose_18.png BODY_PART_INDEXES = { "Head": (16, 14, 0, 15, 17), "Neck": (0, 1), "Shoulder": (2, 5), "Torso": (2, 5, 8, 11), "RArm": (2, 3), "RForearm": (3, 4), "LArm": (5, 6), "LForearm": (6, 7), "RThigh": (8, 9), "RLeg": (9, 10), "LThigh": (11, 12), "LLeg": (12, 13) } BODY_PART_DEFAULT_W_H = { "Head": "256, 256", "Neck": "100, 100", "Shoulder": '', "Torso": "350, 450", "RArm": "128, 256", "RForearm": "128, 256", "LArm": "128, 256", "LForearm": "128, 256", "RThigh": "128, 256", "RLeg": "128, 256", "LThigh": "128, 256", "LLeg": "128, 256" } class SinglePersonProcess: @classmethod def sort_and_get_max_people(s, pose_kps): for idx in range(len(pose_kps)): pose_kps[idx]["people"] = sorted(pose_kps[idx]["people"], key=lambda person:person["pose_keypoints_2d"][0]) return pose_kps, max(len(frame["people"]) for frame in pose_kps) def __init__(self, pose_kps, person_idx=0) -> None: self.width, self.height = pose_kps[0]["canvas_width"], pose_kps[0]["canvas_height"] self.poses = [ self.normalize(pose_frame["people"][person_idx]["pose_keypoints_2d"]) if person_idx < len(pose_frame["people"]) else None for pose_frame in pose_kps ] def normalize(self, pose_kps_2d): n = len(pose_kps_2d) // 3 pose_kps_2d = rearrange(np.array(pose_kps_2d), "(n c) -> n c", n=n, c=3) pose_kps_2d[np.argwhere(pose_kps_2d[:,2]==0), :] = np.iinfo(np.int32).max // 2 #Safe large value pose_kps_2d = pose_kps_2d[:, :2] if is_normalized(pose_kps_2d): pose_kps_2d *= (self.width, self.height) return pose_kps_2d def get_xyxy_bboxes(self, part_name, bbox_size=(128, 256)): width, height = bbox_size xyxy_bboxes = {} for idx, pose in enumerate(self.poses): if pose is None: xyxy_bboxes[idx] = (np.iinfo(np.int32).max // 2,) * 4 continue pts = pose[BODY_PART_INDEXES[part_name], :] #top_left = np.min(pts[:,0]), np.min(pts[:,1]) #bottom_right = np.max(pts[:,0]), np.max(pts[:,1]) #pad_width = np.maximum(width - (bottom_right[0]-top_left[0]), 0) / 2 #pad_height = np.maximum(height - (bottom_right[1]-top_left[1]), 0) / 2 #xyxy_bboxes.append(( # top_left[0] - pad_width, top_left[1] - pad_height, # bottom_right[0] + pad_width, bottom_right[1] + pad_height, #)) x_mid, y_mid = np.mean(pts[:, 0]), np.mean(pts[:, 1]) xyxy_bboxes[idx] = ( x_mid - width/2, y_mid - height/2, x_mid + width/2, y_mid + height/2 ) return xyxy_bboxes class UpperBodyTrackingFromPoseKps: PART_NAMES = ["Head", "Neck", "Shoulder", "Torso", "RArm", "RForearm", "LArm", "LForearm"] @classmethod def INPUT_TYPES(s): return { "required": { "pose_kps": ("POSE_KEYPOINT",), "id_include": ("STRING", {"default": '', "multiline": False}), **{part_name + "_width_height": ("STRING", {"default": BODY_PART_DEFAULT_W_H[part_name], "multiline": False}) for part_name in s.PART_NAMES} } } RETURN_TYPES = ("TRACKING", "STRING") RETURN_NAMES = ("tracking", "prompt") FUNCTION = "convert" CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess" def convert(self, pose_kps, id_include, **parts_width_height): parts_width_height = {part_name.replace("_width_height", ''): value for part_name, value in parts_width_height.items()} enabled_part_names = [part_name for part_name in self.PART_NAMES if len(parts_width_height[part_name].strip())] tracked = {part_name: {} for part_name in enabled_part_names} id_include = id_include.strip() id_include = list(map(int, id_include.split(','))) if len(id_include) else [] prompt_string = '' pose_kps, max_people = SinglePersonProcess.sort_and_get_max_people(pose_kps) for person_idx in range(max_people): if len(id_include) and person_idx not in id_include: continue processor = SinglePersonProcess(pose_kps, person_idx) for part_name in enabled_part_names: bbox_size = tuple(map(int, parts_width_height[part_name].split(','))) 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, "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)", }