diff --git a/openposeeditor/openpose_editor_nodes.py b/openposeeditor/openpose_editor_nodes.py new file mode 100644 index 0000000..0281d38 --- /dev/null +++ b/openposeeditor/openpose_editor_nodes.py @@ -0,0 +1,301 @@ +import json +import torch +import numpy as np +import cv2 +from .util import draw_pose_json, draw_pose + +OpenposeJSON = dict + +class OpenposeEditorNode: + @classmethod + def INPUT_TYPES(s): + return { + "optional": { + "show_body": ("BOOLEAN", {"default": True}), + "show_face": ("BOOLEAN", {"default": True}), + "show_hands": ("BOOLEAN", {"default": True}), + "resolution_x": ("INT", {"default": -1, "min": -1, "max": 12800}), + "use_ground_plane": ("BOOLEAN", {"default": True}), + "pose_marker_size": ("INT", { "default": 4, "min": 0, "max": 100 }), + "face_marker_size": ("INT", { "default": 3, "min": 0, "max": 100 }), + "hand_marker_size": ("INT", { "default": 2, "min": 0, "max": 100 }), + "pelvis_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + "torso_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + "neck_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + # --- 머리 및 눈 관련 스케일 --- + "head_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "eye_distance_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "eye_height": ("FLOAT", { "default": 0.0, "min": -100.0, "max": 100.0, "step": 0.1 }), + "eyebrow_height": ("FLOAT", { "default": 0.0, "min": -100.0, "max": 100.0, "step": 0.1 }), # 눈썹 높이 조절 + "left_eye_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "right_eye_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "left_eyebrow_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "right_eyebrow_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "mouth_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "nose_scale_face": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + "face_shape_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01 }), + # --- + "shoulder_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + "arm_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + "leg_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + "hands_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + "overall_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01 }), + "rotate_angle": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.1}), + "translate_x": ("FLOAT", {"default": 0.0, "min": -10000.0, "max": 10000.0, "step": 0.1}), + "translate_y": ("FLOAT", {"default": 0.0, "min": -10000.0, "max": 10000.0, "step": 0.1}), + "POSE_JSON": ("STRING", {"multiline": True}), + "POSE_KEYPOINT": ("POSE_KEYPOINT",{"default": None}), + "Target_pose_keypoint": ("POSE_KEYPOINT", {"default": None}), + }, + } + + RETURN_NAMES = ("POSE_IMAGE", "POSE_KEYPOINT", "POSE_JSON") + RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT", "STRING") + OUTPUT_NODE = True + FUNCTION = "load_pose" + CATEGORY = "ToyxyzTestNodes" + + def load_pose(self, show_body, show_face, show_hands, resolution_x, use_ground_plane, + pose_marker_size, face_marker_size, hand_marker_size, + pelvis_scale, torso_scale, neck_scale, head_scale, eye_distance_scale, eye_height, eyebrow_height, + left_eye_scale, right_eye_scale, left_eyebrow_scale, right_eyebrow_scale, + mouth_scale, nose_scale_face, face_shape_scale, + shoulder_scale, arm_scale, leg_scale, hands_scale, overall_scale, + rotate_angle, translate_x, translate_y, + POSE_JSON: str, POSE_KEYPOINT=None, Target_pose_keypoint=None) -> tuple[OpenposeJSON]: + + # 내부 함수인 process_pose에 Target_pose_keypoint를 전달하도록 수정 + def process_pose(pose_input_str_list, target_pose_obj=None, rot_angle=0, tr_x=0.0, tr_y=0.0): + pose_imgs, final_keypoints_batch = draw_pose_json( + pose_input_str_list, resolution_x, use_ground_plane, show_body, show_face, show_hands, + pose_marker_size, face_marker_size, hand_marker_size, + pelvis_scale, torso_scale, neck_scale, head_scale, eye_distance_scale, eye_height, eyebrow_height, + left_eye_scale, right_eye_scale, left_eyebrow_scale, right_eyebrow_scale, + mouth_scale, nose_scale_face, face_shape_scale, + shoulder_scale, arm_scale, leg_scale, hands_scale, overall_scale, + rot_angle, tr_x, tr_y, + target_pose_keypoint_obj=target_pose_obj # util.py 함수로 Target_pose_keypoint 전달 + ) + + if not pose_imgs: return None, None, None + + pose_imgs_np = np.array(pose_imgs).astype(np.float32) / 255 + final_json_str = json.dumps(final_keypoints_batch, indent=4) + return torch.from_numpy(pose_imgs_np), final_keypoints_batch, final_json_str + + input_json_str = "" + # 팔 길이 비교를 위해 POSE_KEYPOINT가 우선순위를 갖도록 순서 조정 + if POSE_KEYPOINT is not None: + normalized_json_data = json.dumps(POSE_KEYPOINT, indent=4).replace("'",'"').replace('None','[]') + if not isinstance(POSE_KEYPOINT, list): + input_json_str = f'[{normalized_json_data}]' + else: + input_json_str = normalized_json_data + elif POSE_JSON: + temp_json = POSE_JSON.replace("'",'"').replace('None','[]') + try: + parsed_json = json.loads(temp_json) + input_json_str = f"[{temp_json}]" if not isinstance(parsed_json, list) else temp_json + except json.JSONDecodeError: input_json_str = f"[{temp_json}]" + + if input_json_str: + # process_pose 호출 시 Target_pose_keypoint 객체를 인자로 전달 + image_tensor, keypoint_obj_batch, json_str_batch = process_pose(input_json_str, Target_pose_keypoint, rotate_angle, translate_x, translate_y) + if image_tensor is not None: + return { "ui": {"POSE_JSON": [json_str_batch]}, "result": (image_tensor, keypoint_obj_batch, json_str_batch) } + + W, H = 512, 768 + blank_person = dict(pose_keypoints_2d=[], face_keypoints_2d=[], hand_left_keypoints_2d=[], hand_right_keypoints_2d=[]) + blank_output_keypoints = [{"people": [blank_person], "canvas_width": W, "canvas_height": H}] + W_scaled = resolution_x if resolution_x >= 64 else W + H_scaled = int(H*(W_scaled*1.0/W)) + blank_pose_for_draw = {"people": [blank_person]} + pose_img = [draw_pose(blank_pose_for_draw, H_scaled, W_scaled, pose_marker_size, face_marker_size, hand_marker_size)] + pose_img_np = np.array(pose_img).astype(np.float32) / 255 + return { "ui": {"POSE_JSON": [json.dumps(blank_output_keypoints)]}, "result": (torch.from_numpy(pose_img_np), blank_output_keypoints, json.dumps(blank_output_keypoints)) } + + +class PoseToMaskNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pose_keypoint": ("POSE_KEYPOINT",), + "line_thickness": ("INT", {"default": 30, "min": 1, "max": 1000, "step": 1, "display": "number"}), + "finger_line_thickness": ("INT", {"default": 10, "min": 1, "max": 1000, "step": 1, "display": "number"}), + "torso_thickness": ("INT", {"default": 30, "min": 1, "max": 1000, "step": 1, "display": "number"}), + "head_size": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 1000.0, "step": 0.01, "display": "number"}), # 얼굴 마스크 크기 조절 + "confidence_threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number"}), + } + # "optional": { # 캔버스 크기 오버라이드 옵션 (필요시 유지) + # "override_canvas_width": ("INT", {"default": 0, "min": 0, "max": 8192}), + # "override_canvas_height": ("INT", {"default": 0, "min": 0, "max": 8192}), + # } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("MASK",) + FUNCTION = "create_body_parts_mask" + CATEGORY = "ToyxyzTestNodes" + + # OpenPose 키포인트 인덱스 정의 (참고용) + # Nose:0, Neck:1, RShoulder:2, RElbow:3, RWrist:4, LShoulder:5, LElbow:6, LWrist:7, + # RHip:8, RKnee:9, RAnkle:10, LHip:11, LKnee:12, LAnkle:13 + + # 팔, 다리, 목 연결선 정의 + LIMB_CONNECTIONS = [ + (0, 1), # 코(0) - 목(1) + (2, 3), (3, 4), # 오른쪽 팔 (RShoulder-RElbow, RElbow-RWrist) + (5, 6), (6, 7), # 왼쪽 팔 (LShoulder-LElbow, LElbow-LWrist) + (8, 9), (9, 10), # 오른쪽 다리 (RHip-RKnee, RKnee-RAnkle) + (11, 12), (12, 13) # 왼쪽 다리 (LHip-LKnee, LAnkle) + ] + + # 몸통을 구성하는 4개 꼭짓점 인덱스 (RShoulder, LShoulder, LHip, RHip 순서) + TORSO_POLYGON_INDICES = [2, 5, 11, 8] + + # 손가락 연결선 정의 (OpenPose hand keypoints 기준) + # 손목 (0)에서 시작하여 각 손가락의 마디를 연결 + # 출처: OpenPose GitHub 또는 문서에서 hand keypoint index 참조 + HAND_CONNECTIONS = [ + # 엄지 (Thumb) + (0, 1), (1, 2), (2, 3), (3, 4), + # 검지 (Index Finger) + (0, 5), (5, 6), (6, 7), (7, 8), + # 중지 (Middle Finger) + (0, 9), (9, 10), (10, 11), (11, 12), + # 약지 (Ring Finger) + (0, 13), (13, 14), (14, 15), (15, 16), + # 새끼손가락 (Little Finger) + (0, 17), (17, 18), (18, 19), (19, 20) + ] + + # 얼굴 키포인트 인덱스 정의 + FACE_KEYPOINT_INDICES = list(range(17)) # 0~16번까지의 얼굴 키포인트 사용 + + def _get_point(self, keypoints_list, index, confidence_threshold): + """Helper function to get a valid point's coordinates if confidence is high enough.""" + if not keypoints_list or (index * 3 + 2) >= len(keypoints_list): + return None + x, y, conf = keypoints_list[index * 3], keypoints_list[index * 3 + 1], keypoints_list[index * 3 + 2] + # 유효한 좌표 (OpenPose에서 0,0은 종종 감지 안됨을 의미)이고 신뢰도가 문턱값 이상일 때 + if conf >= confidence_threshold and (x != 0 or y != 0): + return (int(round(x)), int(round(y))) + return None + + def create_body_parts_mask(self, pose_keypoint, line_thickness, finger_line_thickness, torso_thickness, head_size, confidence_threshold, override_canvas_width=0, override_canvas_height=0): + if not pose_keypoint or not isinstance(pose_keypoint, list) or not pose_keypoint[0]: + # 입력이 유효하지 않으면 기본 빈 마스크 반환 + h = override_canvas_height if override_canvas_height > 0 else 768 + w = override_canvas_width if override_canvas_width > 0 else 512 + mask = np.zeros((h, w), dtype=np.float32) + return (torch.from_numpy(mask).unsqueeze(0),) + + frame_data = pose_keypoint[0] # 첫 번째 프레임의 데이터 사용 가정 + + canvas_width = override_canvas_width if override_canvas_width > 0 else frame_data.get("canvas_width", 512) + canvas_height = override_canvas_height if override_canvas_height > 0 else frame_data.get("canvas_height", 768) + + if canvas_width <= 0 or canvas_height <= 0: # 안전장치 + canvas_width = 512 + canvas_height = 768 + + # 최종 마스크는 8비트 단일 채널 + mask_image = np.zeros((canvas_height, canvas_width), dtype=np.uint8) + # 몸통 마스크를 위한 임시 마스크 + torso_mask_temp = np.zeros((canvas_height, canvas_width), dtype=np.uint8) + + + for person_data in frame_data.get("people", []): + body_keypoints = person_data.get("pose_keypoints_2d") + hand_left_keypoints = person_data.get("hand_left_keypoints_2d", []) + hand_right_keypoints = person_data.get("hand_right_keypoints_2d", []) + face_keypoints = person_data.get("face_keypoints_2d", []) + + if not body_keypoints: + continue + + # 1. 몸통 마스크 (채워진 사각형/다각형) - 임시 마스크에 그립니다. + torso_points_for_poly = [] + for idx in self.TORSO_POLYGON_INDICES: + point = self._get_point(body_keypoints, idx, confidence_threshold) + if point: + # 좌표가 캔버스 범위 내에 있는지 확인 + if 0 <= point[0] < canvas_width and 0 <= point[1] < canvas_height: + torso_points_for_poly.append(point) + else: + # 몸통 꼭짓점 중 하나라도 유효하지 않으면 몸통 마스크를 그리지 않음 + torso_points_for_poly = [] # 리스트 비우기 + break + + if len(torso_points_for_poly) == 4: # 4개의 꼭짓점이 모두 유효할 때만 그림 + np_torso_points = np.array([torso_points_for_poly], dtype=np.int32) + cv2.fillConvexPoly(torso_mask_temp, np_torso_points, 255, lineType=cv2.LINE_AA) # 임시 마스크에 흰색(255)으로 채움 + + # 몸통 마스크 확장 (dilate) - 임시 마스크에만 적용 + # torso_thickness 값에 따라 커널 크기 조정 (홀수로 유지) + # torso_thickness가 1보다 작으면 1로 설정하여 최소한의 dilate 적용 + kernel_size_torso = max(1, torso_thickness // 2 * 2 + 1) + kernel_torso = np.ones((kernel_size_torso, kernel_size_torso), np.uint8) + torso_mask_temp = cv2.dilate(torso_mask_temp, kernel_torso, iterations=1) + + # 확장된 몸통 마스크를 최종 마스크에 추가 + mask_image = cv2.bitwise_or(mask_image, torso_mask_temp) + + # 2. 팔, 다리, 목 마스크 (두꺼운 선) - 최종 마스크에 직접 그립니다. + for p1_idx, p2_idx in self.LIMB_CONNECTIONS: + p1 = self._get_point(body_keypoints, p1_idx, confidence_threshold) + p2 = self._get_point(body_keypoints, p2_idx, confidence_threshold) + + if p1 and p2: + # 두 점이 모두 유효하고 캔버스 범위 내에 있을 경우 선 그리기 + if (0 <= p1[0] < canvas_width and 0 <= p1[1] < canvas_height and + 0 <= p2[0] < canvas_width and 0 <= p2[1] < canvas_height): + cv2.line(mask_image, p1, p2, 255, line_thickness, lineType=cv2.LINE_AA) + + # 3. 손가락 마스크 (새로 추가) - 최종 마스크에 직접 그립니다. + for hand_keypoints in [hand_left_keypoints, hand_right_keypoints]: + if not hand_keypoints: + continue + + # 손가락 연결선 그리기 + for p1_idx, p2_idx in self.HAND_CONNECTIONS: + p1 = self._get_point(hand_keypoints, p1_idx, confidence_threshold) + p2 = self._get_point(hand_keypoints, p2_idx, confidence_threshold) + + if p1 and p2: + if (0 <= p1[0] < canvas_width and 0 <= p1[1] < canvas_height and + 0 <= p2[0] < canvas_width and 0 <= p2[1] < canvas_height): + cv2.line(mask_image, p1, p2, 255, finger_line_thickness, lineType=cv2.LINE_AA) + + # 4. 얼굴 마스크 (타원) + face_points = [] + for idx in self.FACE_KEYPOINT_INDICES: + point = self._get_point(face_keypoints, idx, confidence_threshold) + if point: + # 좌표가 캔버스 범위 내에 있는지 확인 + if 0 <= point[0] < canvas_width and 0 <= point[1] < canvas_height: + face_points.append(point) + + if len(face_points) >= 5: # 타원을 그리기 위한 최소 점 개수 (최소 5개의 점이 필요함) + np_face_points = np.array(face_points, dtype=np.int32) + + # 볼록 껍질(Convex Hull) 계산 + hull = cv2.convexHull(np_face_points) + + # 볼록 껍질을 감싸는 최소 크기 타원 계산 + # cv2.fitEllipse는 최소 5개의 점이 필요합니다. + (x, y), (major_axis, minor_axis), angle = cv2.fitEllipse(hull) + + # head_size에 따라 타원 크기 조정 + major_axis *= head_size + minor_axis *= head_size + + # 타원 그리기 + cv2.ellipse(mask_image, (int(x), int(y)), (int(major_axis / 2), int(minor_axis / 2)), int(angle), 0, 360, 255, cv2.FILLED, lineType=cv2.LINE_AA) + + + # NumPy 배열을 PyTorch 텐서로 변환하고 정규화 (H, W) -> (1, H, W) + mask_tensor = torch.from_numpy(mask_image.astype(np.float32) / 255.0).unsqueeze(0) + + return (mask_tensor,) diff --git a/openposeeditor/poseinter.py b/openposeeditor/poseinter.py new file mode 100644 index 0000000..87516d6 --- /dev/null +++ b/openposeeditor/poseinter.py @@ -0,0 +1,823 @@ +import torch +import cv2 +import numpy as np +import copy +from matplotlib.colors import hsv_to_rgb +import json + +DEFAULT_BODY_LIMB_THICKNESS = 6 +DEFAULT_BODY_POINT_RADIUS = 5 +DEFAULT_HAND_LIMB_THICKNESS = 2 +DEFAULT_HAND_POINT_RADIUS = 3 +DEFAULT_FACE_POINT_RADIUS = 2 + +# --- 스켈레톤, 색상, KP 딕셔너리 정의 --- +body_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] +] +face_color = [255, 255, 255] +hand_keypoint_color = [0, 0, 255] +hand_limb_colors = [ + [255,0,0],[255,60,0],[255,120,0],[255,180,0], [180,255,0],[120,255,0],[60,255,0],[0,255,0], + [0,255,60],[0,255,120],[0,255,180],[0,180,255], [0,120,255],[0,60,255],[0,0,255],[60,0,255], + [120,0,255],[180,0,255],[255,0,180],[255,0,120] +] +body_skeleton = [ + [1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [1, 11], + [11, 12], [12, 13], [1, 0], [0, 14], [14, 16], [0, 15], [15, 17] +] +face_skeleton = [] +hand_skeleton = [ + [0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], [10, 11], [11, 12], + [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20] +] + +KP = { + "Nose": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4, + "LShoulder": 5, "LElbow": 6, "LWrist": 7, "MidHip": 8, "RHip": 9, + "RKnee": 10, "RAnkle": 11, "LHip": 12, "LKnee": 13, "LAnkle": 14, + "REye": 15, "LEye": 16, "REar": 17, "LEar": 18, "LBigToe": 19, + "LSmallToe": 20, "LHeel": 21, "RBigToe": 22, "RSmallToe": 23, "RHeel": 24 +} + +def calculate_bone_length(kps, p1_idx, p2_idx): + if kps.shape[0] <= max(p1_idx, p2_idx): return 0.0 + if kps[p1_idx, 2] == 0 or kps[p2_idx, 2] == 0: return 0.0 + p1 = kps[p1_idx, :2] + p2 = kps[p2_idx, :2] + return np.linalg.norm(p1 - p2) + +def get_valid_kps_coords(kps_np, confidence_threshold=0.1): + if kps_np is None or kps_np.ndim != 2 or kps_np.shape[1] != 3 or kps_np.size == 0: return None + valid_points = kps_np[kps_np[:, 2] > confidence_threshold][:, :2] + return valid_points if valid_points.shape[0] > 0 else None + +def get_bounding_box_area_and_center(kps_np, confidence_threshold=0.1): + if kps_np is None or kps_np.ndim != 2 or kps_np.shape[1] != 3 or kps_np.size == 0: + return 0.0, None + valid_points_xy = [] + for i in range(kps_np.shape[0]): + if kps_np[i, 2] > confidence_threshold: + valid_points_xy.append(kps_np[i, :2]) + if not valid_points_xy or len(valid_points_xy) < 1: + return 0.0, None + valid_points_xy = np.array(valid_points_xy) + min_x, min_y = np.min(valid_points_xy, axis=0) + max_x, max_y = np.max(valid_points_xy, axis=0) + width, height = max_x - min_x, max_y - min_y + area = 0.0 + if valid_points_xy.shape[0] >=2 : + area = width * height if width > 1e-6 and height > 1e-6 else 0.0 + center = (np.mean(valid_points_xy[:, 0]), np.mean(valid_points_xy[:, 1])) + return area, center + +def adjust_pose_to_reference_size(source_kps, ref_kps, confidence_threshold=0.1): + if source_kps.size == 0 or ref_kps.size == 0: + return source_kps + + adjusted_kps = source_kps.copy() + + RShoulder, LShoulder, RHip, LHip, Neck = KP["RShoulder"], KP["LShoulder"], KP["RHip"], KP["LHip"], KP["Neck"] + required_indices = [RShoulder, LShoulder, RHip, LHip, Neck] + if not all(idx < source_kps.shape[0] and source_kps[idx, 2] > confidence_threshold for idx in required_indices) or \ + not all(idx < ref_kps.shape[0] and ref_kps[idx, 2] > confidence_threshold for idx in required_indices): + pass + else: + src_shoulder_width = np.linalg.norm(source_kps[LShoulder, :2] - source_kps[RShoulder, :2]) + src_shoulder_center = 0.5 * (source_kps[LShoulder, :2] + source_kps[RShoulder, :2]) + src_hip_center = 0.5 * (source_kps[LHip, :2] + source_kps[RHip, :2]) + src_torso_height = np.linalg.norm(src_shoulder_center - src_hip_center) + ref_shoulder_width = np.linalg.norm(ref_kps[LShoulder, :2] - ref_kps[RShoulder, :2]) + ref_shoulder_center = 0.5 * (ref_kps[LShoulder, :2] + ref_kps[RShoulder, :2]) + ref_hip_center = 0.5 * (ref_kps[LHip, :2] + ref_kps[RHip, :2]) + ref_torso_height = np.linalg.norm(ref_shoulder_center - ref_hip_center) + x_ratio = ref_shoulder_width / src_shoulder_width if src_shoulder_width > 1e-6 else 1.0 + y_ratio = ref_torso_height / src_torso_height if src_torso_height > 1e-6 else 1.0 + neck_pos = adjusted_kps[Neck, :2].copy() + for i in range(adjusted_kps.shape[0]): + if adjusted_kps[i, 2] > 0: + vec_from_neck = adjusted_kps[i, :2] - neck_pos + vec_from_neck[0] *= x_ratio + vec_from_neck[1] *= y_ratio + adjusted_kps[i, :2] = neck_pos + vec_from_neck + + bones_to_adjust = [ + (KP["Neck"], KP["Nose"], [KP["Nose"], KP["REye"], KP["LEye"], KP["REar"], KP["LEar"]]), + (KP["RShoulder"], KP["RElbow"], [KP["RElbow"], KP["RWrist"]]), + (KP["RElbow"], KP["RWrist"], [KP["RWrist"]]), + (KP["LShoulder"], KP["LElbow"], [KP["LElbow"], KP["LWrist"]]), + (KP["LElbow"], KP["LWrist"], [KP["LWrist"]]), + (KP["RHip"], KP["RKnee"], [KP["RKnee"], KP["RAnkle"], KP["RBigToe"], KP["RSmallToe"], KP["RHeel"]]), + (KP["RKnee"], KP["RAnkle"], [KP["RAnkle"], KP["RBigToe"], KP["RSmallToe"], KP["RHeel"]]), + (KP["LHip"], KP["LKnee"], [KP["LKnee"], KP["LAnkle"], KP["LBigToe"], KP["LSmallToe"], KP["LHeel"]]), + (KP["LKnee"], KP["LAnkle"], [KP["LAnkle"], KP["LBigToe"], KP["LSmallToe"], KP["LHeel"]]), + ] + for parent_idx, child_idx, children_indices in bones_to_adjust: + if max(parent_idx, child_idx) >= adjusted_kps.shape[0] or max(parent_idx, child_idx) >= ref_kps.shape[0]: continue + if adjusted_kps[parent_idx, 2] < confidence_threshold or adjusted_kps[child_idx, 2] < confidence_threshold or \ + ref_kps[parent_idx, 2] < confidence_threshold or ref_kps[child_idx, 2] < confidence_threshold: + continue + len_source = calculate_bone_length(adjusted_kps, parent_idx, child_idx) + len_ref = calculate_bone_length(ref_kps, parent_idx, child_idx) + if len_source == 0 or len_ref == 0: continue + ratio = len_ref / len_source + if abs(1.0 - ratio) < 0.01: continue + parent_pos = adjusted_kps[parent_idx, :2] + child_pos_old = adjusted_kps[child_idx, :2] + vector = child_pos_old - parent_pos + vector_new = vector * ratio + child_pos_new = parent_pos + vector_new + offset = child_pos_new - child_pos_old + for idx in children_indices: + if idx < adjusted_kps.shape[0] and adjusted_kps[idx, 2] > 0: + adjusted_kps[idx, :2] += offset + + if Neck < adjusted_kps.shape[0] and Neck < ref_kps.shape[0] and \ + adjusted_kps[Neck, 2] > confidence_threshold and ref_kps[Neck, 2] > confidence_threshold: + final_offset = ref_kps[Neck, :2] - adjusted_kps[Neck, :2] + for i in range(adjusted_kps.shape[0]): + if adjusted_kps[i, 2] > 0: + adjusted_kps[i, :2] += final_offset + + head_kp_indices = [KP["Nose"], KP["REye"], KP["LEye"], KP["REar"], KP["LEar"]] + ref_head_points = np.array([ref_kps[i] for i in head_kp_indices if i < ref_kps.shape[0] and ref_kps[i, 2] > confidence_threshold]) + adj_head_points = np.array([adjusted_kps[i] for i in head_kp_indices if i < adjusted_kps.shape[0] and adjusted_kps[i, 2] > confidence_threshold]) + if ref_head_points.shape[0] >= 2 and adj_head_points.shape[0] >= 2: + ref_head_area, _ = get_bounding_box_area_and_center(ref_head_points, confidence_threshold) + adj_head_area, adj_head_center = get_bounding_box_area_and_center(adj_head_points, confidence_threshold) + if adj_head_area > 1e-6 and ref_head_area > 1e-6 and adj_head_center is not None: + scale_factor_head = np.sqrt(ref_head_area / adj_head_area) + if abs(1.0 - scale_factor_head) > 0.01: + center_x, center_y = adj_head_center + for kp_idx in head_kp_indices: + if kp_idx < adjusted_kps.shape[0] and adjusted_kps[kp_idx, 2] > confidence_threshold: + x, y, _ = adjusted_kps[kp_idx] + adjusted_kps[kp_idx, 0] = center_x + (x - center_x) * scale_factor_head + adjusted_kps[kp_idx, 1] = center_y + (y - center_y) * scale_factor_head + return adjusted_kps + +def adjust_face_keypoints_size(full_face_kps_np, scale_factor, center_xy, confidence_threshold=0.1): + adjusted_kps = full_face_kps_np.copy() + if center_xy is None : return adjusted_kps + center_x, center_y = center_xy + for i in range(adjusted_kps.shape[0]): + x, y, conf = adjusted_kps[i] + if conf > confidence_threshold: + adjusted_kps[i, :2] = (center_x + (x - center_x) * scale_factor, center_y + (y - center_y) * scale_factor) + return adjusted_kps + +def adjust_face_to_maintain_relative_offset(orig_target_body_kps, orig_target_face_kps, adjusted_body_kps, adjusted_face_kps, confidence_threshold=0.1): + FACE_NOSE_INDEX = 8 + if not (orig_target_body_kps.shape[0] > KP["Nose"] and orig_target_body_kps[KP["Nose"], 2] >= confidence_threshold and + orig_target_face_kps.shape[0] > FACE_NOSE_INDEX and orig_target_face_kps[FACE_NOSE_INDEX, 2] >= confidence_threshold and + adjusted_body_kps.shape[0] > KP["Nose"] and adjusted_body_kps[KP["Nose"], 2] >= confidence_threshold and + adjusted_face_kps.shape[0] > FACE_NOSE_INDEX and adjusted_face_kps[FACE_NOSE_INDEX, 2] >= confidence_threshold): + return adjusted_face_kps + + orig_body_nose_pos = orig_target_body_kps[KP["Nose"], :2] + orig_face_nose_pos = orig_target_face_kps[FACE_NOSE_INDEX, :2] + original_offset = orig_face_nose_pos - orig_body_nose_pos + + adjusted_body_nose_pos = adjusted_body_kps[KP["Nose"], :2] + current_face_nose_pos = adjusted_face_kps[FACE_NOSE_INDEX, :2] + + desired_face_nose_pos = adjusted_body_nose_pos + original_offset + translation_vector = desired_face_nose_pos - current_face_nose_pos + + final_face_kps = adjusted_face_kps.copy() + for i in range(final_face_kps.shape[0]): + if final_face_kps[i, 2] > confidence_threshold: + final_face_kps[i, :2] += translation_vector + return final_face_kps + +def calculate_hand_intrinsic_properties(hand_kps_np, confidence_threshold=0.1): + if hand_kps_np is None or hand_kps_np.size == 0: return None + hand_kp0_abs = None + if hand_kps_np.shape[0] > 0 and hand_kps_np[0, 2] > confidence_threshold: + hand_kp0_abs = hand_kps_np[0, :2].copy() + + valid_hand_coords_xy = get_valid_kps_coords(hand_kps_np, confidence_threshold) + scale = 0.0 + if valid_hand_coords_xy is not None and valid_hand_coords_xy.shape[0] >= 2: + min_x, min_y = np.min(valid_hand_coords_xy, axis=0) + max_x, max_y = np.max(valid_hand_coords_xy, axis=0) + width, height = max_x - min_x, max_y - min_y + scale = np.sqrt(width**2 + height**2) if width > 1e-6 and height > 1e-6 else 0.0 + return {'scale': scale, 'kp0_abs': hand_kp0_abs} + +def transform_hand_final(target_hand_kps_np_orig, target_body_wrist_pos_xy_adj, ref_hand_kps_np_orig, ref_body_wrist_pos_xy_orig, confidence_threshold=0.1): + ref_props = calculate_hand_intrinsic_properties(ref_hand_kps_np_orig, confidence_threshold) + target_orig_props = calculate_hand_intrinsic_properties(target_hand_kps_np_orig, confidence_threshold) + + if not all([ref_props, target_orig_props, + ref_props.get('kp0_abs') is not None, target_orig_props.get('kp0_abs') is not None, + ref_body_wrist_pos_xy_orig is not None, target_body_wrist_pos_xy_adj is not None]): + return target_hand_kps_np_orig.flatten().tolist() if target_hand_kps_np_orig is not None and target_hand_kps_np_orig.size > 0 else [] + + ref_hand_kp0_pos, ref_scale = ref_props['kp0_abs'], ref_props['scale'] + target_orig_hand_kp0_pos, target_orig_scale = target_orig_props['kp0_abs'], target_orig_props['scale'] + + ref_offset_bodywrist_to_handkp0 = ref_hand_kp0_pos - ref_body_wrist_pos_xy_orig + scale_factor = ref_scale / target_orig_scale if target_orig_scale > 1e-6 else 1.0 + + scaled_target_hand_kps = target_hand_kps_np_orig.copy() + pivot_for_scaling = target_orig_hand_kp0_pos.copy() + + for i in range(scaled_target_hand_kps.shape[0]): + if scaled_target_hand_kps[i, 2] > confidence_threshold: + vec_from_pivot = scaled_target_hand_kps[i, :2] - pivot_for_scaling + scaled_target_hand_kps[i, :2] = pivot_for_scaling + (vec_from_pivot * scale_factor) + + if scaled_target_hand_kps.shape[0] == 0 or scaled_target_hand_kps[0, 2] <= confidence_threshold: + return scaled_target_hand_kps.flatten().tolist() + + current_abs_pos_of_scaled_target_hand_kp0 = scaled_target_hand_kps[0, :2] + desired_abs_pos_for_target_hand_kp0 = target_body_wrist_pos_xy_adj + ref_offset_bodywrist_to_handkp0 + translation_vector = desired_abs_pos_for_target_hand_kp0 - current_abs_pos_of_scaled_target_hand_kp0 + + for i in range(scaled_target_hand_kps.shape[0]): + if scaled_target_hand_kps[i, 2] > confidence_threshold: + scaled_target_hand_kps[i, :2] += translation_vector + + return scaled_target_hand_kps.flatten().tolist() + +def draw_keypoints_and_skeleton(image, keypoints_data, skeleton_connections, + colors_config, + limb_thickness, + point_radius, + confidence_threshold=0.1, + is_body=False, is_face=False, is_hand=False, + hand_edges_count=0): + if not keypoints_data or len(keypoints_data) % 3 != 0: return + tri_tuples = [keypoints_data[i:i + 3] for i in range(0, len(keypoints_data), 3)] + + if skeleton_connections: + for i, (joint_idx_a, joint_idx_b) in enumerate(skeleton_connections): + if joint_idx_a >= len(tri_tuples) or joint_idx_b >= len(tri_tuples): continue + + a_x_f, a_y_f, a_confidence = tri_tuples[joint_idx_a] + b_x_f, b_y_f, b_confidence = tri_tuples[joint_idx_b] + + if a_confidence >= confidence_threshold and b_confidence >= confidence_threshold: + a_x, a_y = int(round(a_x_f)), int(round(a_y_f)) + b_x, b_y = int(round(b_x_f)), int(round(b_y_f)) + + current_limb_color = None + if is_body: + current_limb_color = tuple(colors_config[i % len(colors_config)]) + elif is_hand: + if hand_edges_count > 0: + rgb_color = hsv_to_rgb([i / float(hand_edges_count), 1.0, 1.0]) + current_limb_color = tuple((np.array(rgb_color) * 255).astype(np.uint8).tolist()) + else: + current_limb_color = tuple(colors_config[i % len(colors_config)]) + elif is_face: + if skeleton_connections: + current_limb_color = tuple(colors_config) + + if current_limb_color is not None: + if is_body: + center_x, center_y = (a_x + b_x) // 2, (a_y + b_y) // 2 + length = np.linalg.norm(np.array([a_x, a_y]) - np.array([b_x, b_y])) + if length >= 1: + angle_rad = np.arctan2(b_y - a_y, b_x - a_x) + angle_deg = np.degrees(angle_rad) + ellipse_major_axis = max(1, int(length / 2)) + ellipse_minor_axis = max(1, int(limb_thickness / 2)) + axes = (ellipse_major_axis, ellipse_minor_axis) + polygon_points = cv2.ellipse2Poly((center_x, center_y), axes, int(angle_deg), 0, 360, 10) + cv2.fillConvexPoly(image, polygon_points, current_limb_color) + else: + cv2.line(image, (a_x, a_y), (b_x, b_y), current_limb_color, limb_thickness) + + for i, (x_f, y_f, confidence) in enumerate(tri_tuples): + if confidence >= confidence_threshold: + current_point_color = None + if is_body: + current_point_color = body_colors[i % len(body_colors)] + elif is_hand: + current_point_color = hand_keypoint_color + elif is_face: + current_point_color = face_color + + if current_point_color: + cv2.circle(image, (int(round(x_f)), int(round(y_f))), point_radius, tuple(current_point_color), -1) + +def gen_skeleton_with_face_hands(pose_keypoints_2d, face_keypoints_2d, hand_left_keypoints_2d, hand_right_keypoints_2d, + canvas_width, canvas_height, landmarkType, confidence_threshold=0.1): + image = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8) + def scale_keypoints(keypoints, target_w, target_h, input_is_normalized): + if not keypoints or len(keypoints) % 3 != 0 : return [] + scaled = [] + for i in range(0, len(keypoints), 3): + x, y, conf = keypoints[i:i+3] + scaled.extend([x * target_w, y * target_h, conf] if input_is_normalized else [x, y, conf]) + return scaled + + input_normalized = (landmarkType == "OpenPose") + + scaled_pose = scale_keypoints(pose_keypoints_2d, canvas_width, canvas_height, input_normalized) + scaled_face = scale_keypoints(face_keypoints_2d, canvas_width, canvas_height, input_normalized) + scaled_hand_left = scale_keypoints(hand_left_keypoints_2d, canvas_width, canvas_height, input_normalized) + scaled_hand_right = scale_keypoints(hand_right_keypoints_2d, canvas_width, canvas_height, input_normalized) + + draw_keypoints_and_skeleton(image, scaled_pose, body_skeleton, body_colors, + DEFAULT_BODY_LIMB_THICKNESS, DEFAULT_BODY_POINT_RADIUS, + confidence_threshold, is_body=True) + + if scaled_face: + draw_keypoints_and_skeleton(image, scaled_face, face_skeleton, face_color, + 0, DEFAULT_FACE_POINT_RADIUS, + confidence_threshold, is_face=True) + + if scaled_hand_left: + draw_keypoints_and_skeleton(image, scaled_hand_left, hand_skeleton, hand_limb_colors, + DEFAULT_HAND_LIMB_THICKNESS, DEFAULT_HAND_POINT_RADIUS, + confidence_threshold, is_hand=True, hand_edges_count=len(hand_skeleton)) + if scaled_hand_right: + draw_keypoints_and_skeleton(image, scaled_hand_right, hand_skeleton, hand_limb_colors, + DEFAULT_HAND_LIMB_THICKNESS, DEFAULT_HAND_POINT_RADIUS, + confidence_threshold, is_hand=True, hand_edges_count=len(hand_skeleton)) + return image + +def transform_all_keypoints(keypoints_1, keypoints_2, frames, interpolation="linear"): + def interpolate_keypoint_set(kp1, kp2, num_frames, interp_method): + kp1 = kp1 if kp1 is not None else [] + kp2 = kp2 if kp2 is not None else [] + + if not kp1 and not kp2: return [[] for _ in range(num_frames)] + + len_kp1 = len(kp1) + len_kp2 = len(kp2) + + if len_kp1 == 0 and len_kp2 > 0: + kp1 = [0.0] * len_kp2 + elif len_kp2 == 0 and len_kp1 > 0: + kp2 = [0.0] * len_kp1 + elif len_kp1 != len_kp2 : + print(f"Warning: Keypoint list length mismatch. kp1 len: {len_kp1}, kp2 len: {len_kp2}. Interpolation might be unreliable.") + max_len = max(len_kp1, len_kp2) + if max_len % 3 != 0 : + print(f"Error: Max keypoint list length {max_len} is not a multiple of 3. Returning empty interpolation.") + return [[] for _ in range(num_frames)] + while len(kp1) < max_len: kp1.extend([0.0, 0.0, 0.0]) + while len(kp2) < max_len: kp2.extend([0.0, 0.0, 0.0]) + + if not kp1 and not kp2: return [[] for _ in range(num_frames)] + + num_kps1 = len(kp1) // 3 + num_kps2 = len(kp2) // 3 + + if num_kps1 != num_kps2: + print(f"Critical Error: Mismatch in number of keypoints after padding. KPs1: {num_kps1}, KPs2: {num_kps2}") + return [[] for _ in range(num_frames)] + if num_kps1 == 0 : return [[] for _ in range(num_frames)] + + tri_tuples_1 = [kp1[i:i + 3] for i in range(0, len(kp1), 3)] + tri_tuples_2 = [kp2[i:i + 3] for i in range(0, len(kp2), 3)] + + keypoints_sequence = [] + for j in range(num_frames): + interpolated_kps_for_frame = [] + t = j / float(num_frames - 1) if num_frames > 1 else 0.0 + + if interp_method == "ease-in": interp_factor = t * t + elif interp_method == "ease-out": interp_factor = 1 - (1 - t) * (1 - t) + elif interp_method == "ease-in-out": + interp_factor = 4 * t * t * t if t < 0.5 else 1.0 - pow(-2 * t + 2, 3) / 2 + else: interp_factor = t + + for i in range(num_kps1): + x1, y1, c1 = tri_tuples_1[i]; x2, y2, c2 = tri_tuples_2[i] + new_x, new_y, new_c = 0.0, 0.0, 0.0 + + if c1 > 0 and c2 > 0: + new_x = x1 + (x2 - x1) * interp_factor + new_y = y1 + (y2 - y1) * interp_factor + new_c = c1 + (c2 - c1) * interp_factor + elif c1 > 0: + new_x, new_y = x1, y1 + new_c = c1 * (1.0 - interp_factor) + elif c2 > 0: + new_x, new_y = x2, y2 + new_c = c2 * interp_factor + interpolated_kps_for_frame.extend([new_x, new_y, new_c]) + keypoints_sequence.append(interpolated_kps_for_frame) + return keypoints_sequence + + parts = ['pose', 'face', 'hand_left', 'hand_right'] + sequences = {} + for part in parts: + kp1_part = keypoints_1.get(f'{part}_keypoints_2d', []) + kp2_part = keypoints_2.get(f'{part}_keypoints_2d', []) + sequences[part] = interpolate_keypoint_set(kp1_part, kp2_part, frames, interpolation) + + combined_sequence = [] + for i in range(frames): + combined_frame_data = {} + valid_frame = False + for part in parts: + if i < len(sequences[part]) and sequences[part][i]: + combined_frame_data[f'{part}_keypoints_2d'] = sequences[part][i] + valid_frame = True + else: + combined_frame_data[f'{part}_keypoints_2d'] = [] + if valid_frame: + combined_sequence.append(combined_frame_data) + return combined_sequence + +def apply_confidence_threshold(keypoints_list, threshold): + if not keypoints_list: + return [] + filtered_kps = [] + for i in range(0, len(keypoints_list), 3): + x, y, c = keypoints_list[i:i+3] + if c < threshold: + filtered_kps.extend([0.0, 0.0, 0.0]) + else: + filtered_kps.extend([x, y, c]) + return filtered_kps + +class Pose_Inter: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + interpolation_methods = ["linear", "ease-in", "ease-out", "ease-in-out"] + return { + "required": { + "pose_from": ("POSE_KEYPOINT", ), "pose_to": ("POSE_KEYPOINT", ), + "interpolate_frames": ("INT", {"default": 12, "min": 1, "max": 99999, "step": 1}), + "interpolation": (interpolation_methods, {"default": "linear"}), + "confidence_threshold": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}), + "adjust_body_shape": ("BOOLEAN", {"default": False}), + "landmarkType": (["OpenPose", "DWPose"], {"default": "DWPose"}), + "include_face": ("BOOLEAN", {"default": True}), + "include_hands": ("BOOLEAN", {"default": True}), + "pick_frame": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}), # Allow negative for pick_frame list items + }, + } + + RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT",) + RETURN_NAMES = ("image", "pose_keypoint",) + FUNCTION = "run" + CATEGORY = "ToyxyzTestNodes" + + def run(self, pose_from, pose_to, interpolate_frames, interpolation, confidence_threshold, landmarkType, include_face, include_hands, adjust_body_shape, pick_frame): + if not pose_from or not pose_to: + raise ValueError("Input 'pose_from' or 'pose_to' data is empty.") + + pose_from_list = pose_from if isinstance(pose_from, list) else [pose_from] + pose_to_list = pose_to if isinstance(pose_to, list) else [pose_to] + + if not pose_from_list or not pose_to_list: + raise ValueError("Input 'pose_from' or 'pose_to' list is empty after ensuring it's a list.") + + if len(pose_from_list) != len(pose_to_list): + raise ValueError(f"Batch size mismatch: 'pose_from' has {len(pose_from_list)} items, 'pose_to' has {len(pose_to_list)} items. They must be equal.") + + batch_size = len(pose_from_list) + + pick_rules_for_each_item_in_batch = [] + if isinstance(pick_frame, list): + # If pick_frame is a list, it applies per batch item if batch_size > 1 + # Or, if batch_size is 1, this list applies to that single item. + if batch_size > 1 and len(pick_frame) != batch_size: + raise ValueError(f"'pick_frame' list has {len(pick_frame)} items, but batch size is {batch_size}. They must be equal if 'pick_frame' is a list for multi-item batch processing where each list item corresponds to a batch item.") + for i in range(batch_size): + if batch_size == 1: + # If batch is 1, the entire pick_frame list is the rule for this one item + pick_rules_for_each_item_in_batch.append(pick_frame) + else: + # If batch > 1, then pick_frame[i] is the rule for pose_from_list[i] + # Ensure it's a list, even if it's a single int from the input pick_frame list + rule = pick_frame[i] + pick_rules_for_each_item_in_batch.append([rule] if isinstance(rule, int) else rule) + elif isinstance(pick_frame, int): + # If pick_frame is a single int, this int (wrapped in a list) becomes the rule for all batch items. + pick_rules_for_each_item_in_batch = [[pick_frame]] * batch_size + else: + raise TypeError("'pick_frame' must be an INT or a LIST of INTs/LISTs.") + + + final_output_images = [] + final_output_poses = [] + + default_canvas_width = 512 + default_canvas_height = 512 + if batch_size > 0 and pose_from_list[0] and isinstance(pose_from_list[0], dict): + default_canvas_width = pose_from_list[0].get("canvas_width", default_canvas_width) + default_canvas_height = pose_from_list[0].get("canvas_height", default_canvas_height) + + + for i in range(batch_size): + current_pose_from_dict = pose_from_list[i] + current_pose_to_dict = pose_to_list[i] + current_pick_rule_list_for_item = pick_rules_for_each_item_in_batch[i] + + if not isinstance(current_pose_from_dict, dict) or "people" not in current_pose_from_dict or not current_pose_from_dict["people"]: + print(f"Warning: Invalid or no people data in 'pose_from' for batch item {i}. Skipping.") + continue + if not isinstance(current_pose_to_dict, dict) or "people" not in current_pose_to_dict or not current_pose_to_dict["people"]: + print(f"Warning: Invalid or no people data in 'pose_to' for batch item {i}. Skipping.") + continue + + person_from = current_pose_from_dict["people"][0] + person_to = current_pose_to_dict["people"][0] + + keypoints_from_current_pair = { + 'pose_keypoints_2d': person_from.get("pose_keypoints_2d", []), + 'face_keypoints_2d': person_from.get("face_keypoints_2d", []) if include_face else [], + 'hand_left_keypoints_2d': person_from.get("hand_left_keypoints_2d", []) if include_hands else [], + 'hand_right_keypoints_2d': person_from.get("hand_right_keypoints_2d", []) if include_hands else [] + } + keypoints_to_current_pair = { + 'pose_keypoints_2d': person_to.get("pose_keypoints_2d", []), + 'face_keypoints_2d': person_to.get("face_keypoints_2d", []) if include_face else [], + 'hand_left_keypoints_2d': person_to.get("hand_left_keypoints_2d", []) if include_hands else [], + 'hand_right_keypoints_2d': person_to.get("hand_right_keypoints_2d", []) if include_hands else [] + } + + original_person_to_face_kps_current = person_to.get("face_keypoints_2d", []) if include_face else [] + original_person_to_hand_left_kps_current = person_to.get("hand_left_keypoints_2d", []) if include_hands else [] + original_person_to_hand_right_kps_current = person_to.get("hand_right_keypoints_2d", []) if include_hands else [] + + kps_from_np_body = np.array(keypoints_from_current_pair['pose_keypoints_2d']).reshape(-1, 3) if keypoints_from_current_pair['pose_keypoints_2d'] else np.array([]) + kps_to_np_body_for_adjustment = np.array(keypoints_to_current_pair['pose_keypoints_2d']).reshape(-1, 3) if keypoints_to_current_pair['pose_keypoints_2d'] else np.array([]) + kps_to_np_body_final_for_interp = kps_to_np_body_for_adjustment.copy() + + if adjust_body_shape: + temp_kps_to_body = kps_to_np_body_for_adjustment.copy() + if kps_from_np_body.size > 0 and temp_kps_to_body.size > 0: + adjusted_body_kps_intermediate = adjust_pose_to_reference_size(temp_kps_to_body, kps_from_np_body, confidence_threshold) + keypoints_to_current_pair['pose_keypoints_2d'] = adjusted_body_kps_intermediate.flatten().tolist() + kps_to_np_body_final_for_interp = adjusted_body_kps_intermediate + + if include_face: + face_kps_from_np = np.array(keypoints_from_current_pair['face_keypoints_2d']).reshape(-1, 3) if keypoints_from_current_pair['face_keypoints_2d'] else np.array([]) + face_kps_to_np_orig = np.array(original_person_to_face_kps_current).reshape(-1, 3) if original_person_to_face_kps_current else np.array([]) + + if face_kps_from_np.size > 0 and face_kps_to_np_orig.size > 0: + scaled_face_kps_to = face_kps_to_np_orig.copy() + area_from, _ = get_bounding_box_area_and_center(face_kps_from_np, confidence_threshold) + area_to_orig, center_to_face_orig = get_bounding_box_area_and_center(face_kps_to_np_orig, confidence_threshold) + if area_to_orig > 1e-6 and area_from > 1e-6 and center_to_face_orig is not None: + scale_factor = np.sqrt(area_from / area_to_orig) + scaled_face_kps_to = adjust_face_keypoints_size(face_kps_to_np_orig, scale_factor, center_to_face_orig, confidence_threshold) + + final_adjusted_face_kps = adjust_face_to_maintain_relative_offset( + kps_to_np_body_for_adjustment, + face_kps_to_np_orig, + kps_to_np_body_final_for_interp, + scaled_face_kps_to, + confidence_threshold + ) + keypoints_to_current_pair['face_keypoints_2d'] = final_adjusted_face_kps.flatten().tolist() + + if include_hands: + for hand_type in ["left", "right"]: + wrist_kp_name = "LWrist" if hand_type == "left" else "RWrist" + ref_body_wrist_pos_xy_orig, target_body_wrist_pos_xy_adj = None, None + + if kps_from_np_body.size > 0 and KP[wrist_kp_name] < kps_from_np_body.shape[0] and kps_from_np_body[KP[wrist_kp_name], 2] > confidence_threshold: + ref_body_wrist_pos_xy_orig = kps_from_np_body[KP[wrist_kp_name], :2] + + if kps_to_np_body_final_for_interp.size > 0 and KP[wrist_kp_name] < kps_to_np_body_final_for_interp.shape[0] and kps_to_np_body_final_for_interp[KP[wrist_kp_name], 2] > confidence_threshold: + target_body_wrist_pos_xy_adj = kps_to_np_body_final_for_interp[KP[wrist_kp_name], :2] + + ref_hand_kps_list = keypoints_from_current_pair.get(f'hand_{hand_type}_keypoints_2d', []) + ref_hand_kps_np_orig = np.array(ref_hand_kps_list).reshape(-1, 3) if ref_hand_kps_list and len(ref_hand_kps_list) % 3 == 0 else np.array([]) + + target_hand_kps_list_orig_current = original_person_to_hand_left_kps_current if hand_type == "left" else original_person_to_hand_right_kps_current + target_hand_kps_np_orig_current = np.array(target_hand_kps_list_orig_current).reshape(-1, 3) if target_hand_kps_list_orig_current and len(target_hand_kps_list_orig_current) % 3 == 0 else np.array([]) + + if ref_body_wrist_pos_xy_orig is not None and target_body_wrist_pos_xy_adj is not None and ref_hand_kps_np_orig.size > 0 and target_hand_kps_np_orig_current.size > 0: + adjusted_hand_kps_list = transform_hand_final( + target_hand_kps_np_orig_current, + target_body_wrist_pos_xy_adj, + ref_hand_kps_np_orig, + ref_body_wrist_pos_xy_orig, + confidence_threshold) + keypoints_to_current_pair[f'hand_{hand_type}_keypoints_2d'] = adjusted_hand_kps_list + + interpolated_sequence_for_current_pair = transform_all_keypoints( + keypoints_from_current_pair, + keypoints_to_current_pair, + interpolate_frames, + interpolation + ) + + frames_to_render_this_item = [] + if not interpolated_sequence_for_current_pair: + print(f"Warning: Interpolation failed for batch item {i}. No frames to pick.") + else: + num_available_frames = len(interpolated_sequence_for_current_pair) + if num_available_frames == 0: + print(f"Warning: Interpolation resulted in zero frames for batch item {i}.") + else: + for pick_value in current_pick_rule_list_for_item: + if not isinstance(pick_value, int): + print(f"Warning: Invalid non-integer pick_value '{pick_value}' for batch item {i}. Skipping this pick_value.") + continue + + target_frame_num_1_based = pick_value + + if target_frame_num_1_based == 0: + if interpolate_frames > 0 : + frames_to_render_this_item.extend(interpolated_sequence_for_current_pair) + else: + if target_frame_num_1_based < 1: + target_frame_num_1_based = 1 # Clamp to 1st frame + if target_frame_num_1_based > num_available_frames: + target_frame_num_1_based = num_available_frames # Clamp to last frame + + frames_to_render_this_item.append(interpolated_sequence_for_current_pair[target_frame_num_1_based - 1]) + + # 중복 프레임 제거 (선택적: 만약 [0, 1] 같은 rule로 인해 중복이 생기는 것을 방지하고 싶다면) + # 이 경우, 추가된 순서가 중요하지 않다면 set으로 변환 후 list로 다시 만들 수 있지만, + # 여기서는 사용자가 명시적으로 여러 번 같은 프레임을 요청할 수도 있으므로 중복 제거 안 함. + + canvas_width_current = current_pose_from_dict.get("canvas_width", default_canvas_width) + canvas_height_current = current_pose_from_dict.get("canvas_height", default_canvas_height) + + for frame_data in frames_to_render_this_item: # 이미 선택/조정된 프레임 데이터 목록 + pose_output_for_this_frame = copy.deepcopy(current_pose_from_dict) + + pose_kps_final = apply_confidence_threshold(frame_data.get('pose_keypoints_2d', []), confidence_threshold) + face_kps_final = apply_confidence_threshold(frame_data.get('face_keypoints_2d', []), confidence_threshold) + hand_left_kps_final = apply_confidence_threshold(frame_data.get('hand_left_keypoints_2d', []), confidence_threshold) + hand_right_kps_final = apply_confidence_threshold(frame_data.get('hand_right_keypoints_2d', []), confidence_threshold) + + if "people" in pose_output_for_this_frame and pose_output_for_this_frame["people"]: + pose_output_for_this_frame["people"][0]["pose_keypoints_2d"] = pose_kps_final + pose_output_for_this_frame["people"][0]["face_keypoints_2d"] = face_kps_final + pose_output_for_this_frame["people"][0]["hand_left_keypoints_2d"] = hand_left_kps_final + pose_output_for_this_frame["people"][0]["hand_right_keypoints_2d"] = hand_right_kps_final + + pose_output_for_this_frame["canvas_width"] = canvas_width_current + pose_output_for_this_frame["canvas_height"] = canvas_height_current + + final_output_poses.append(pose_output_for_this_frame) + + image_np = gen_skeleton_with_face_hands( + pose_kps_final, + face_kps_final, + hand_left_kps_final, + hand_right_kps_final, + canvas_width_current, + canvas_height_current, + landmarkType, + confidence_threshold + ) + image_tensor = torch.from_numpy(image_np.astype(np.float32) / 255.0) + final_output_images.append(image_tensor) + + if not final_output_images: + print("Warning: No images were generated across all batch items. Returning a single black image.") + black_image_np = np.zeros((default_canvas_height, default_canvas_width, 3), dtype=np.float32) + return (torch.from_numpy(black_image_np).unsqueeze(0), []) + + return (torch.stack(final_output_images), final_output_poses) + +class PoseKeypointToCoordStr: # + def __init__(self): # + pass # + + @classmethod + def INPUT_TYPES(cls): # + return { # + "required": { # + "pose_keypoint": ("POSE_KEYPOINT",), # + "enable_body": ("BOOLEAN", {"default": True, "label_on": "Body Enabled", "label_off": "Body Disabled"}), + "enable_face": ("BOOLEAN", {"default": True, "label_on": "Face Enabled", "label_off": "Face Disabled"}), + "enable_hand": ("BOOLEAN", {"default": True, "label_on": "Hands Enabled", "label_off": "Hands Disabled"}), + "enable_extra_points": ("BOOLEAN", {"default": False, "label_on": "Extra Body Points Enabled", "label_off": "Extra Body Points Disabled"}), + "num_extra_points_per_bone": ("INT", {"default": 5, "min": 0, "max": 1000, "step": 1, "label": "Extra Points Per Bone"}), + } + } + + RETURN_TYPES = ("STRING",) # ComfyUI에서 문자열 리스트를 담는 단일 슬롯으로 처리될 수 있음 + RETURN_NAMES = ("coord_str",) # + FUNCTION = "convert_to_coord_str" # + CATEGORY = "ToyxyzTestNodes" # + + def convert_to_coord_str(self, pose_keypoint, enable_body, enable_face, enable_hand, enable_extra_points, num_extra_points_per_bone): # num_extra_points_per_bone 파라미터 추가 + if not pose_keypoint: + return (["[]"],) + + pose_keypoint_list = pose_keypoint if isinstance(pose_keypoint, list) else [pose_keypoint] + + if not pose_keypoint_list: + return (["[]"],) + + all_poses_kps_extracted = [] + max_kps_count_overall = 0 + + # NUM_EXTRA_POINTS_PER_BONE 상수를 제거하고 입력 파라미터 사용 + + for pose_data_idx, pose_data in enumerate(pose_keypoint_list): + current_pose_all_coords_dicts_for_frame = [] + if pose_data and "people" in pose_data and pose_data["people"]: + person_data = pose_data["people"][0] + + # Body points 처리 + if enable_body: + body_keypoints_flat = person_data.get("pose_keypoints_2d", []) + if body_keypoints_flat and len(body_keypoints_flat) > 0: + body_kps_triplets = [] + for i in range(0, len(body_keypoints_flat), 3): + x = int(body_keypoints_flat[i]) + y = int(body_keypoints_flat[i+1]) + c = body_keypoints_flat[i+2] + body_kps_triplets.append({"x": x, "y": y, "c": c}) + current_pose_all_coords_dicts_for_frame.append({"x": x, "y": y}) + + # enable_extra_points가 True이고, 사용자가 지정한 추가 포인트 수가 0보다 클 경우 + if enable_extra_points and num_extra_points_per_bone > 0: + # body_skeleton은 전역 변수로 가정 (코드 상단에 정의된 것을 사용) + for p1_idx, p2_idx in body_skeleton: + if p1_idx < len(body_kps_triplets) and p2_idx < len(body_kps_triplets): + p1 = body_kps_triplets[p1_idx] + p2 = body_kps_triplets[p2_idx] + + if p1["c"] > 0 and p2["c"] > 0: # 두 원본 포인트의 신뢰도가 유효할 때만 + for j in range(1, num_extra_points_per_bone + 1): + ratio = j / float(num_extra_points_per_bone + 1) + extra_x = int(p1["x"] + (p2["x"] - p1["x"]) * ratio) + extra_y = int(p1["y"] + (p2["y"] - p1["y"]) * ratio) + current_pose_all_coords_dicts_for_frame.append({"x": extra_x, "y": extra_y}) + + # Face points 처리 + if enable_face: + face_keypoints_flat = person_data.get("face_keypoints_2d", []) + if face_keypoints_flat and len(face_keypoints_flat) > 0: + for i in range(0, len(face_keypoints_flat), 3): + x = int(face_keypoints_flat[i]) + y = int(face_keypoints_flat[i+1]) + current_pose_all_coords_dicts_for_frame.append({"x": x, "y": y}) + + # Hand points 처리 (left and right) + if enable_hand: + for hand_type_key in ["hand_left_keypoints_2d", "hand_right_keypoints_2d"]: + hand_keypoints_flat = person_data.get(hand_type_key, []) + if hand_keypoints_flat and len(hand_keypoints_flat) > 0: + for i in range(0, len(hand_keypoints_flat), 3): + x = int(hand_keypoints_flat[i]) + y = int(hand_keypoints_flat[i+1]) + current_pose_all_coords_dicts_for_frame.append({"x": x, "y": y}) + + all_poses_kps_extracted.append(current_pose_all_coords_dicts_for_frame) + if len(current_pose_all_coords_dicts_for_frame) > max_kps_count_overall: + max_kps_count_overall = len(current_pose_all_coords_dicts_for_frame) + + if max_kps_count_overall == 0: + return (["[]"],) + + output_coord_groups_json_str = [] + for i in range(max_kps_count_overall): + coords_for_this_track = [] + for single_pose_kps_list in all_poses_kps_extracted: + if i < len(single_pose_kps_list): + coords_for_this_track.append(single_pose_kps_list[i]) + else: + if coords_for_this_track: + coords_for_this_track.append({"x": coords_for_this_track[-1]["x"], "y": coords_for_this_track[-1]["y"]}) + else: + coords_for_this_track.append({"x": 0, "y": 0}) + + output_coord_groups_json_str.append(json.dumps(coords_for_this_track)) + + return (output_coord_groups_json_str,) + +class JoinPose: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pose_keypoint_1": ("POSE_KEYPOINT",), + "pose_keypoint_2": ("POSE_KEYPOINT",), + } + } + + RETURN_TYPES = ("POSE_KEYPOINT",) + RETURN_NAMES = ("pose_keypoint",) + FUNCTION = "join_poses" + CATEGORY = "ToyxyzTestNodes" + + def join_poses(self, pose_keypoint_1, pose_keypoint_2): + + # 입력이 None일 경우 빈 리스트로 처리하여 오류 방지 + list_1 = pose_keypoint_1 if pose_keypoint_1 is not None else [] + list_2 = pose_keypoint_2 if pose_keypoint_2 is not None else [] + + # 두 리스트를 순서대로 합침 + joined_list = list(list_1) + list(list_2) + + return (joined_list,) + diff --git a/openposeeditor/util.py b/openposeeditor/util.py new file mode 100644 index 0000000..73e119c --- /dev/null +++ b/openposeeditor/util.py @@ -0,0 +1,651 @@ +import math +import json +import numpy as np +import matplotlib +import cv2 +from comfy.utils import ProgressBar + +eps = 0.01 + +def scale(point, scale_factor, pivot): + if not isinstance(point, np.ndarray): point = np.array(point) + if not isinstance(pivot, np.ndarray): pivot = np.array(pivot) + return pivot + (point - pivot) * scale_factor + +def draw_pose_json(pose_json_str, resolution_x, use_ground_plane, show_body, show_face, show_hands, + pose_marker_size, face_marker_size, hand_marker_size, + pelvis_scale, torso_scale, neck_scale, head_scale, eye_distance_scale, eye_height, eyebrow_height, + left_eye_scale, right_eye_scale, left_eyebrow_scale, right_eyebrow_scale, + mouth_scale, nose_scale_face, face_shape_scale, + shoulder_scale, arm_scale, leg_scale, hands_scale, overall_scale, + rotate_angle, translate_x, translate_y, + target_pose_keypoint_obj=None): + + # 최종적으로 적용될 스케일 값을 초기화 + final_hands_scale = hands_scale + final_torso_scale = torso_scale + final_head_scale = head_scale + final_neck_scale = neck_scale + final_pelvis_scale = pelvis_scale + final_shoulder_scale = shoulder_scale + final_arm_scale = arm_scale + final_leg_scale = leg_scale + + if target_pose_keypoint_obj and pose_json_str: + try: + source_pose_obj = json.loads(pose_json_str) + + # --- 공용 헬퍼 함수 정의 --- + def get_point(kps_list, index): + if index * 3 + 2 >= len(kps_list) or kps_list[index * 3 + 2] == 0: + return None + return np.array([kps_list[index * 3], kps_list[index * 3 + 1]]) + + def calculate_limb_length(kps_list, p1_idx, p2_idx): + p1 = get_point(kps_list, p1_idx) + p2 = get_point(kps_list, p2_idx) + if p1 is not None and p2 is not None: + return np.linalg.norm(p1 - p2) + return 0.0 + + # --- 팔 길이 계산 --- + def get_max_arm_length(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + keypoints = pose_obj[0]['people'][0].get('pose_keypoints_2d', []) + if not keypoints: return 0.0 + right_arm_len = calculate_limb_length(keypoints, 2, 3) + calculate_limb_length(keypoints, 3, 4) + left_arm_len = calculate_limb_length(keypoints, 5, 6) + calculate_limb_length(keypoints, 6, 7) + return max(left_arm_len, right_arm_len) + except (IndexError, TypeError): return 0.0 + + target_arm_len = get_max_arm_length(target_pose_keypoint_obj) + source_arm_len = get_max_arm_length(source_pose_obj) + + if source_arm_len > 0 and target_arm_len > 0: + final_arm_scale = arm_scale * (target_arm_len / source_arm_len) + + # --- 다리 길이 계산 --- + def get_max_leg_length(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + keypoints = pose_obj[0]['people'][0].get('pose_keypoints_2d', []) + if not keypoints: return 0.0 + right_leg_len = calculate_limb_length(keypoints, 8, 9) + calculate_limb_length(keypoints, 9, 10) + left_leg_len = calculate_limb_length(keypoints, 11, 12) + calculate_limb_length(keypoints, 12, 13) + return max(left_leg_len, right_leg_len) + except (IndexError, TypeError): return 0.0 + + target_leg_len = get_max_leg_length(target_pose_keypoint_obj) + source_leg_len = get_max_leg_length(source_pose_obj) + + if source_leg_len > 0 and target_leg_len > 0: + final_leg_scale = leg_scale * (target_leg_len / source_leg_len) + + # --- 어깨 너비 계산 --- + def get_shoulder_width(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + keypoints = pose_obj[0]['people'][0].get('pose_keypoints_2d', []) + if not keypoints: return 0.0 + width = calculate_limb_length(keypoints, 2, 5) + return width + except (IndexError, TypeError): return 0.0 + + target_shoulder_width = get_shoulder_width(target_pose_keypoint_obj) + source_shoulder_width = get_shoulder_width(source_pose_obj) + + if source_shoulder_width > 0 and target_shoulder_width > 0: + final_shoulder_scale = shoulder_scale * (target_shoulder_width / source_shoulder_width) + + # --- 골반 너비 계산 --- + def get_pelvis_width(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + keypoints = pose_obj[0]['people'][0].get('pose_keypoints_2d', []) + if not keypoints: return 0.0 + width = calculate_limb_length(keypoints, 8, 11) + return width + except (IndexError, TypeError): return 0.0 + + target_pelvis_width = get_pelvis_width(target_pose_keypoint_obj) + source_pelvis_width = get_pelvis_width(source_pose_obj) + + if source_pelvis_width > 0 and target_pelvis_width > 0: + final_pelvis_scale = pelvis_scale * (target_pelvis_width / source_pelvis_width) + + # --- 목 길이 계산 --- + def get_neck_length(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + keypoints = pose_obj[0]['people'][0].get('pose_keypoints_2d', []) + if not keypoints: return 0.0 + length = calculate_limb_length(keypoints, 1, 0) + return length + except (IndexError, TypeError): return 0.0 + + target_neck_length = get_neck_length(target_pose_keypoint_obj) + source_neck_length = get_neck_length(source_pose_obj) + + if source_neck_length > 0 and target_neck_length > 0: + final_neck_scale = neck_scale * (target_neck_length / source_neck_length) + + # --- 머리 크기 계산 --- + def get_head_size(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + keypoints = pose_obj[0]['people'][0].get('pose_keypoints_2d', []) + if not keypoints: return 0.0 + + head_indices = [0, 14, 15, 16, 17] + valid_points = [p for i in head_indices if (p := get_point(keypoints, i)) is not None] + + if len(valid_points) < 3: return 0.0 + + points_for_hull = np.array(valid_points, dtype=np.float32).reshape((-1, 1, 2)) + hull = cv2.convexHull(points_for_hull) + area = cv2.contourArea(hull) + return area + except (IndexError, TypeError): return 0.0 + + target_head_size = get_head_size(target_pose_keypoint_obj) + source_head_size = get_head_size(source_pose_obj) + + if source_head_size > 0 and target_head_size > 0: + size_ratio = math.sqrt(target_head_size / source_head_size) + final_head_scale = head_scale * size_ratio + + # --- 몸통 길이 계산 --- + def get_torso_length(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + keypoints = pose_obj[0]['people'][0].get('pose_keypoints_2d', []) + if not keypoints: return 0.0 + + right_hip = get_point(keypoints, 8) + left_hip = get_point(keypoints, 11) + neck = get_point(keypoints, 1) + + if right_hip is None or left_hip is None or neck is None: return 0.0 + + hip_midpoint = (right_hip + left_hip) / 2.0 + length = np.linalg.norm(neck - hip_midpoint) + return length + except (IndexError, TypeError): return 0.0 + + target_torso_length = get_torso_length(target_pose_keypoint_obj) + source_torso_length = get_torso_length(source_pose_obj) + + if source_torso_length > 0 and target_torso_length > 0: + final_torso_scale = torso_scale * (target_torso_length / source_torso_length) + + # --- 손 크기 계산 (새로 추가된 부분) --- + def get_hand_area(hand_kps_list): + if not hand_kps_list: return 0.0 + valid_points = [] + for i in range(0, len(hand_kps_list), 3): + if hand_kps_list[i+2] > 0: + valid_points.append([hand_kps_list[i], hand_kps_list[i+1]]) + if len(valid_points) < 3: return 0.0 + points_for_hull = np.array(valid_points, dtype=np.float32).reshape((-1, 1, 2)) + hull = cv2.convexHull(points_for_hull) + return cv2.contourArea(hull) + + def get_max_hand_size(pose_obj): + try: + if not isinstance(pose_obj, list) or not pose_obj or 'people' not in pose_obj[0] or not pose_obj[0]['people']: return 0.0 + person = pose_obj[0]['people'][0] + left_hand_kps = person.get('hand_left_keypoints_2d', []) + right_hand_kps = person.get('hand_right_keypoints_2d', []) + left_area = get_hand_area(left_hand_kps) + right_area = get_hand_area(right_hand_kps) + return max(left_area, right_area) + except(IndexError, TypeError): return 0.0 + + target_hand_size = get_max_hand_size(target_pose_keypoint_obj) + source_hand_size = get_max_hand_size(source_pose_obj) + + if source_hand_size > 0 and target_hand_size > 0: + size_ratio = math.sqrt(target_hand_size / source_hand_size) + final_hands_scale = hands_scale * size_ratio + + + except (json.JSONDecodeError, IndexError, TypeError): + # 에러 발생 시 원래 값 유지 + final_hands_scale = hands_scale + final_torso_scale = torso_scale + final_head_scale = head_scale + final_neck_scale = neck_scale + final_pelvis_scale = pelvis_scale + final_shoulder_scale = shoulder_scale + final_arm_scale = arm_scale + final_leg_scale = leg_scale + + pose_imgs = [] + all_frames_keypoints_output = [] + + if pose_json_str: + images_data_list = json.loads(pose_json_str) + if not isinstance(images_data_list, list): images_data_list = [images_data_list] + + pbar = ProgressBar(len(images_data_list)) + + KP = { + "Nose": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4, + "LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9, + "RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "REye": 14, + "LEye": 15, "REar": 16, "LEar": 17 + } + + FACE_KP_GROUPS_INDICES = { + "Left_Eye": [42, 43, 44, 45, 46, 47, 69], + "Right_Eye": [36, 37, 38, 39, 40, 41, 68], + "Left_Eyebrow": [22, 23, 24, 25, 26], + "Right_Eyebrow": [17, 18, 19, 20, 21], + "Mouth": list(range(48, 68)), + "Nose_Face": list(range(27, 36)), + "Face_Shape": list(range(0, 17)) + } + + INDIVIDUAL_FACE_SCALES = { + "Left_Eye": left_eye_scale, "Right_Eye": right_eye_scale, + "Left_Eyebrow": left_eyebrow_scale, "Right_Eyebrow": right_eyebrow_scale, + "Mouth": mouth_scale, "Nose_Face": nose_scale_face, + "Face_Shape": face_shape_scale + } + + BODY_HEAD_PARTS = {KP["REye"], KP["LEye"], KP["REar"], KP["LEar"]} + R_LEG_INDICES = {KP["RKnee"], KP["RAnkle"]} + L_LEG_INDICES = {KP["LKnee"], KP["LAnkle"]} + FEET_INDICES = {KP["RAnkle"], KP["LAnkle"]} + + for image_data in images_data_list: + if 'people' not in image_data or not image_data['people']: + pbar.update(1); continue + + figures = image_data['people'] + H = image_data['canvas_height'] + W = image_data['canvas_width'] + + current_image_people_data_for_output = [] + all_scaled_candidates_for_drawing, all_scaled_faces_for_drawing, all_scaled_hands_for_drawing = [], [], [] + final_subset_for_drawing = [[]] + + for fig_idx, figure in enumerate(figures): + body_raw, face_raw, lhand_raw, rhand_raw = [figure.get(k, []) for k in ['pose_keypoints_2d', 'face_keypoints_2d', 'hand_left_keypoints_2d', 'hand_right_keypoints_2d']] + + if not body_raw or len(body_raw) < (KP["LEar"] + 1) * 3: continue + + initial_candidate = np.array([body_raw[i:i+2] for i in range(0, len(body_raw), 3)]) + confidence_scores_body = [body_raw[i*3+2] for i in range(len(initial_candidate))] + scaled_candidate_np = initial_candidate.copy() + + r_hip_orig, l_hip_orig = initial_candidate[KP["RHip"]], initial_candidate[KP["LHip"]] + hip_center_orig = (r_hip_orig + l_hip_orig) / 2 + neck_orig, r_shoulder_orig, l_shoulder_orig, nose_orig = [initial_candidate[KP[k]] for k in ["Neck", "RShoulder", "LShoulder", "Nose"]] + lwrist_orig, rwrist_orig = initial_candidate[KP["LWrist"]], initial_candidate[KP["RWrist"]] + + r_hip_final = scale(r_hip_orig, final_pelvis_scale, hip_center_orig) + l_hip_final = scale(l_hip_orig, final_pelvis_scale, hip_center_orig) + scaled_candidate_np[KP["RHip"]], scaled_candidate_np[KP["LHip"]] = r_hip_final, l_hip_final + + hip_center_final = (r_hip_final + l_hip_final) / 2 + neck_final = scale(neck_orig, final_torso_scale, hip_center_final) + scaled_candidate_np[KP["Neck"]] = neck_final + + scaled_candidate_np[KP["RShoulder"]] = neck_final + (r_shoulder_orig - neck_orig) * final_shoulder_scale + scaled_candidate_np[KP["LShoulder"]] = neck_final + (l_shoulder_orig - neck_orig) * final_shoulder_scale + + r_shoulder_final, l_shoulder_final = scaled_candidate_np[KP["RShoulder"]], scaled_candidate_np[KP["LShoulder"]] + + for i in [KP["RElbow"], KP["RWrist"]]: scaled_candidate_np[i] = r_shoulder_final + (initial_candidate[i] - r_shoulder_orig) * final_arm_scale + for i in [KP["LElbow"], KP["LWrist"]]: scaled_candidate_np[i] = l_shoulder_final + (initial_candidate[i] - l_shoulder_orig) * final_arm_scale + + for i in R_LEG_INDICES: scaled_candidate_np[i] = r_hip_final + (initial_candidate[i] - r_hip_orig) * final_leg_scale + for i in L_LEG_INDICES: scaled_candidate_np[i] = l_hip_final + (initial_candidate[i] - l_hip_orig) * final_leg_scale + + nose_body_final = neck_final + (nose_orig - neck_orig) * final_neck_scale + scaled_candidate_np[KP["Nose"]] = nose_body_final + + effective_nose_translation = nose_body_final - nose_orig + for i in BODY_HEAD_PARTS: + part_moved_with_nose = initial_candidate[i] + effective_nose_translation + scaled_candidate_np[i] = scale(part_moved_with_nose, final_head_scale, nose_body_final) + + face_points_scaled_current_fig = [] + if face_raw: + face_points_orig = [np.array(face_raw[i:i+2]) for i in range(0, len(face_raw), 3)] + num_face_points = len(face_points_orig) + + face_points_positioned = [p + effective_nose_translation for p in face_points_orig] + face_points_after_global_head_scale = [scale(p, final_head_scale, nose_body_final) for p in face_points_positioned] + face_points_scaled_current_fig = list(face_points_after_global_head_scale) + + reye_pos_after_head_scale = scale(initial_candidate[KP["REye"]] + effective_nose_translation, final_head_scale, nose_body_final) + leye_pos_after_head_scale = scale(initial_candidate[KP["LEye"]] + effective_nose_translation, final_head_scale, nose_body_final) + eye_center = (reye_pos_after_head_scale + leye_pos_after_head_scale) / 2 + reye_pos_after_dist_scale = scale(reye_pos_after_head_scale, eye_distance_scale, eye_center) + leye_pos_after_dist_scale = scale(leye_pos_after_head_scale, eye_distance_scale, eye_center) + right_dist_translation = reye_pos_after_dist_scale - reye_pos_after_head_scale + left_dist_translation = leye_pos_after_dist_scale - leye_pos_after_head_scale + + eye_height_offset = np.array([0.0, 0.0]) + eyebrow_height_offset = np.array([0.0, 0.0]) + direction_vector = nose_body_final - neck_final + norm_direction = np.linalg.norm(direction_vector) + if norm_direction > eps: + unit_direction = direction_vector / norm_direction + if abs(eye_height) > eps: eye_height_offset = unit_direction * eye_height + if abs(eyebrow_height) > eps: eyebrow_height_offset = unit_direction * eyebrow_height + + group_translations = { + "Right_Eye": right_dist_translation + eye_height_offset, + "Left_Eye": left_dist_translation + eye_height_offset, + "Right_Eyebrow": right_dist_translation + eyebrow_height_offset, + "Left_Eyebrow": left_dist_translation + eyebrow_height_offset, + } + + scaled_candidate_np[KP["REye"]] = reye_pos_after_dist_scale + eye_height_offset + scaled_candidate_np[KP["LEye"]] = leye_pos_after_dist_scale + eye_height_offset + + for group_name, indices in FACE_KP_GROUPS_INDICES.items(): + group_scale_modifier = INDIVIDUAL_FACE_SCALES.get(group_name, 1.0) + valid_indices = [idx for idx in indices if idx < num_face_points] + if not valid_indices: continue + + points_after_head_scale = [face_points_after_global_head_scale[idx] for idx in valid_indices] + + if group_name in group_translations: + points_after_translation = [p + group_translations[group_name] for p in points_after_head_scale] + else: + points_after_translation = points_after_head_scale + + if abs(group_scale_modifier - 1.0) > eps: + if group_name == "Face_Shape": + pivot = nose_body_final + direction_vector = neck_final - nose_body_final + norm_direction = np.linalg.norm(direction_vector) + if norm_direction > eps: + unit_direction = direction_vector / norm_direction + final_points = [] + for p in points_after_translation: + point_vector = p - pivot + proj_length = np.dot(point_vector, unit_direction) + parallel_component = proj_length * unit_direction + perpendicular_component = point_vector - parallel_component + scaled_parallel_component = parallel_component * group_scale_modifier + new_point = pivot + scaled_parallel_component + perpendicular_component + final_points.append(new_point) + else: + final_points = points_after_translation + else: + pivot = np.mean(points_after_translation, axis=0) + final_points = [scale(p, group_scale_modifier, pivot) for p in points_after_translation] + else: + final_points = points_after_translation + + for i, idx in enumerate(valid_indices): + face_points_scaled_current_fig[idx] = final_points[i] + + lwrist_final_calc, rwrist_final_calc = scaled_candidate_np[KP["LWrist"]], scaled_candidate_np[KP["RWrist"]] + + # hands_scale 대신 계산된 final_hands_scale을 사용하도록 수정 + lhand_scaled_current_fig = [(scale(np.array(lhand_raw[i:i+2]), final_hands_scale, lwrist_orig) + (lwrist_final_calc - lwrist_orig)) if lhand_raw[i+2] > 0 else np.array([0.0, 0.0]) for i in range(0, len(lhand_raw), 3)] if lhand_raw else [] + rhand_scaled_current_fig = [(scale(np.array(rhand_raw[i:i+2]), final_hands_scale, rwrist_orig) + (rwrist_final_calc - rwrist_orig)) if rhand_raw[i+2] > 0 else np.array([0.0, 0.0]) for i in range(0, len(rhand_raw), 3)] if rhand_raw else [] + + scales_to_check = [leg_scale, torso_scale, overall_scale, pelvis_scale, head_scale] + is_scaling_active = any(abs(s - 1.0) > 0.001 for s in scales_to_check) + + candidate_list_current_fig_np = scaled_candidate_np + face_list_current_fig_np = np.array(face_points_scaled_current_fig) if face_points_scaled_current_fig else np.array([]) + lhand_list_current_fig_np = np.array(lhand_scaled_current_fig) if lhand_scaled_current_fig else np.array([]) + rhand_list_current_fig_np = np.array(rhand_scaled_current_fig) if rhand_scaled_current_fig else np.array([]) + + if use_ground_plane and is_scaling_active: + ground_y_coord = H + orig_feet_coords = [initial_candidate[i] for i in FEET_INDICES if i < len(initial_candidate)] + orig_lowest_y = max(p[1] for p in orig_feet_coords) if orig_feet_coords else H + orig_dist_to_ground = ground_y_coord - orig_lowest_y + + feet_coords_for_overall_pivot = [candidate_list_current_fig_np[i] for i in FEET_INDICES if i < len(candidate_list_current_fig_np)] + + if feet_coords_for_overall_pivot: + feet_pos_pivot = np.mean(feet_coords_for_overall_pivot, axis=0) + candidate_list_current_fig_np = np.array([scale(p, overall_scale, feet_pos_pivot) for p in candidate_list_current_fig_np]) + if face_list_current_fig_np.size > 0: face_list_current_fig_np = np.array([scale(p, overall_scale, feet_pos_pivot) for p in face_list_current_fig_np]) + if lhand_list_current_fig_np.size > 0: lhand_list_current_fig_np = np.array([scale(p, overall_scale, feet_pos_pivot) if np.sum(np.abs(p)) > eps else p for p in lhand_list_current_fig_np]) + if rhand_list_current_fig_np.size > 0: rhand_list_current_fig_np = np.array([scale(p, overall_scale, feet_pos_pivot) if np.sum(np.abs(p)) > eps else p for p in rhand_list_current_fig_np]) + + final_feet_coords = [candidate_list_current_fig_np[i] for i in FEET_INDICES if i < len(candidate_list_current_fig_np)] + if final_feet_coords: + final_lowest_y = max(p[1] for p in final_feet_coords) + desired_final_y = ground_y_coord - orig_dist_to_ground + vertical_translation = desired_final_y - final_lowest_y + + candidate_list_current_fig_np = candidate_list_current_fig_np + np.array([0, vertical_translation]) + if face_list_current_fig_np.size > 0: face_list_current_fig_np = face_list_current_fig_np + np.array([0, vertical_translation]) + if lhand_list_current_fig_np.size > 0: lhand_list_current_fig_np = lhand_list_current_fig_np + np.array([0, vertical_translation]) + if rhand_list_current_fig_np.size > 0: rhand_list_current_fig_np = rhand_list_current_fig_np + np.array([0, vertical_translation]) + else: + center_pivot = [W * 0.5, H * 0.5] + candidate_list_current_fig_np = np.array([scale(p, overall_scale, center_pivot) for p in candidate_list_current_fig_np]) + if face_list_current_fig_np.size > 0: face_list_current_fig_np = np.array([scale(p, overall_scale, center_pivot) for p in face_list_current_fig_np]) + if lhand_list_current_fig_np.size > 0: lhand_list_current_fig_np = np.array([scale(p, overall_scale, center_pivot) if np.sum(np.abs(p)) > eps else p for p in lhand_list_current_fig_np]) + if rhand_list_current_fig_np.size > 0: rhand_list_current_fig_np = np.array([scale(p, overall_scale, center_pivot) if np.sum(np.abs(p)) > eps else p for p in rhand_list_current_fig_np]) + + # Rotation Logic (after all scaling, before translation) + if abs(rotate_angle) > eps: # Only rotate if angle is significant + all_points_for_rotation_center = [] + if candidate_list_current_fig_np.size > 0: + # 유효한 body 포인트만 중심 계산에 사용 (confidence 기반으로 필터링하는 것이 더 정확할 수 있으나, 여기서는 모든 점 사용) + all_points_for_rotation_center.extend(candidate_list_current_fig_np.tolist()) + if face_list_current_fig_np.size > 0: + all_points_for_rotation_center.extend(face_list_current_fig_np.tolist()) + + # 손 포인트 중 [0,0]이 아닌 유효한 포인트만 중심 계산에 사용 + if lhand_list_current_fig_np.size > 0: + valid_lhand_points = [p.tolist() for p in lhand_list_current_fig_np if np.sum(np.abs(p)) > eps] + if valid_lhand_points: + all_points_for_rotation_center.extend(valid_lhand_points) + if rhand_list_current_fig_np.size > 0: + valid_rhand_points = [p.tolist() for p in rhand_list_current_fig_np if np.sum(np.abs(p)) > eps] + if valid_rhand_points: + all_points_for_rotation_center.extend(valid_rhand_points) + + if all_points_for_rotation_center: + points_for_center_np = np.array(all_points_for_rotation_center) + center_x = np.mean(points_for_center_np[:, 0]) + center_y = np.mean(points_for_center_np[:, 1]) + + angle_rad = math.radians(rotate_angle) + cos_a = math.cos(angle_rad) + sin_a = math.sin(angle_rad) + + def apply_rotation_to_points(points_np, cx, cy, c_angle, s_angle): + if points_np.size == 0: + return points_np + + # 회전 적용할 포인트만 선택 (예: [0,0] 제외는 여기서 처리 안함, 모든 점 동일하게 회전) + # 원본 포인트를 복사하여 사용 + rotated_points = points_np.copy() + + # 중심점으로 이동 + translated_x = rotated_points[:, 0] - cx + translated_y = rotated_points[:, 1] - cy + + # 회전 + rotated_x = translated_x * c_angle - translated_y * s_angle + rotated_y = translated_x * s_angle + translated_y * c_angle + + # 다시 원래 위치로 이동 (중심점 기준) + rotated_points[:, 0] = rotated_x + cx + rotated_points[:, 1] = rotated_y + cy + return rotated_points + + if candidate_list_current_fig_np.size > 0: + candidate_list_current_fig_np = apply_rotation_to_points(candidate_list_current_fig_np, center_x, center_y, cos_a, sin_a) + if face_list_current_fig_np.size > 0: + face_list_current_fig_np = apply_rotation_to_points(face_list_current_fig_np, center_x, center_y, cos_a, sin_a) + if lhand_list_current_fig_np.size > 0: + # [0,0] 점들도 회전 중심에 대해 상대적으로 회전됨 + lhand_list_current_fig_np = apply_rotation_to_points(lhand_list_current_fig_np, center_x, center_y, cos_a, sin_a) + if rhand_list_current_fig_np.size > 0: + # [0,0] 점들도 회전 중심에 대해 상대적으로 회전됨 + rhand_list_current_fig_np = apply_rotation_to_points(rhand_list_current_fig_np, center_x, center_y, cos_a, sin_a) + + + if abs(translate_x) > eps or abs(translate_y) > eps: # 실제로 이동이 필요한 경우에만 연산 + translation_vector = np.array([translate_x, translate_y], dtype=np.float32) + + if candidate_list_current_fig_np.size > 0: + candidate_list_current_fig_np = candidate_list_current_fig_np + translation_vector + + if face_list_current_fig_np.size > 0: + face_list_current_fig_np = face_list_current_fig_np + translation_vector + + if lhand_list_current_fig_np.size > 0: + lhand_list_current_fig_np = lhand_list_current_fig_np + translation_vector + + if rhand_list_current_fig_np.size > 0: + rhand_list_current_fig_np = rhand_list_current_fig_np + translation_vector + + + body_kps_out_current_fig = [item for i, p in enumerate(candidate_list_current_fig_np) for item in [p[0], p[1], confidence_scores_body[i]]] + face_kps_out_current_fig = [item for p in face_list_current_fig_np for item in [p[0], p[1], 1.0]] if face_list_current_fig_np.size > 0 else [] + + original_lhand_confidences = [lhand_raw[i+2] for i in range(0, len(lhand_raw), 3)] if lhand_raw else [] + original_rhand_confidences = [rhand_raw[i+2] for i in range(0, len(rhand_raw), 3)] if rhand_raw else [] + + lhand_kps_out_current_fig = [item for i, p in enumerate(lhand_list_current_fig_np) for item in [p[0], p[1], original_lhand_confidences[i]]] if lhand_list_current_fig_np.size > 0 else [] + rhand_kps_out_current_fig = [item for i, p in enumerate(rhand_list_current_fig_np) for item in [p[0], p[1], original_rhand_confidences[i]]] if rhand_list_current_fig_np.size > 0 else [] + + current_image_people_data_for_output.append({ + "pose_keypoints_2d": body_kps_out_current_fig, "face_keypoints_2d": face_kps_out_current_fig, + "hand_left_keypoints_2d": lhand_kps_out_current_fig, "hand_right_keypoints_2d": rhand_kps_out_current_fig, + }) + + all_scaled_candidates_for_drawing.extend(candidate_list_current_fig_np.tolist()) + if face_list_current_fig_np.size > 0: all_scaled_faces_for_drawing.extend(face_list_current_fig_np.tolist()) + if lhand_list_current_fig_np.size > 0: all_scaled_hands_for_drawing.append(lhand_list_current_fig_np.tolist()) + if rhand_list_current_fig_np.size > 0: all_scaled_hands_for_drawing.append(rhand_list_current_fig_np.tolist()) + + if fig_idx == 0 and not final_subset_for_drawing[0]: + final_subset_for_drawing[0].extend([i if body_raw[i*3+2]>0 else -1 for i in range(len(candidate_list_current_fig_np))]) + else: + prev_candidate_count = len(all_scaled_candidates_for_drawing) - len(candidate_list_current_fig_np) + final_subset_for_drawing.append([prev_candidate_count+i if body_raw[i*3+2]>0 else -1 for i in range(len(candidate_list_current_fig_np))]) + + current_frame_keypoint_object = { "people": current_image_people_data_for_output, "canvas_width": W, "canvas_height": H } + all_frames_keypoints_output.append(current_frame_keypoint_object) + + candidate_norm, faces_norm = all_scaled_candidates_for_drawing, all_scaled_faces_for_drawing + hands_norm_for_drawing = all_scaled_hands_for_drawing + + if candidate_norm: + candidate_np_norm = np.array(candidate_norm).astype(float); candidate_np_norm[...,0] /= float(W); candidate_np_norm[...,1] /= float(H) + candidate_norm = candidate_np_norm.tolist() + if faces_norm: + faces_np_norm = np.array(faces_norm).astype(float); + if faces_np_norm.size > 0: faces_np_norm[...,0] /= float(W); faces_np_norm[...,1] /= float(H) + faces_norm = faces_np_norm.tolist() + + hands_final_norm_for_drawing = [] + if hands_norm_for_drawing: + for hand_kps_list in hands_norm_for_drawing: + current_normalized_hand = [] + for point_list in hand_kps_list: + if not isinstance(point_list, (list, np.ndarray)) or len(point_list) != 2: continue + norm_point = np.array(point_list).astype(float) + if norm_point[0] > eps or norm_point[1] > eps: + norm_point[0] /= float(W) + norm_point[1] /= float(H) + current_normalized_hand.append(norm_point.tolist()) + if current_normalized_hand : hands_final_norm_for_drawing.append(current_normalized_hand) + + bodies = dict(candidate=candidate_norm, subset=final_subset_for_drawing) + original_face_exists = any(fig.get('face_keypoints_2d') for fig in figures) + original_lhand_exists = any(fig.get('hand_left_keypoints_2d') for fig in figures) + original_rhand_exists = any(fig.get('hand_right_keypoints_2d') for fig in figures) + + pose = dict( + bodies=bodies if show_body else {'candidate':[], 'subset':[]}, + faces=faces_norm if show_face and original_face_exists else [], + hands=hands_final_norm_for_drawing if show_hands and (original_lhand_exists or original_rhand_exists) else [] + ) + W_scaled = resolution_x if resolution_x >= 64 else W + H_scaled = int(H*(W_scaled*1.0/W)) + pose_imgs.append(draw_pose(pose, H_scaled, W_scaled, pose_marker_size, face_marker_size, hand_marker_size)) + pbar.update(1) + + return pose_imgs, all_frames_keypoints_output + +def draw_pose(pose, H, W, pose_marker_size, face_marker_size, hand_marker_size): + canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8) + body_render_info = pose.get('bodies', {}) + candidate = body_render_info.get('candidate', []) + subset = body_render_info.get('subset', []) + faces_data = pose.get('faces', []) + hands_data = pose.get('hands', []) + + if candidate and subset and np.array(candidate).size > 0 : canvas = draw_bodypose(canvas, np.array(candidate), np.array(subset), pose_marker_size) + if hands_data and np.array(hands_data).size > 0 : canvas = draw_handpose(canvas, hands_data, hand_marker_size) + if faces_data and np.array(faces_data).size > 0 : canvas = draw_facepose(canvas, faces_data, face_marker_size) + return canvas + +def draw_bodypose(canvas, candidate, subset, pose_marker_size): + H, W, C = canvas.shape + limbSeq = [[1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [1, 11], [11, 12], [12, 13], [1, 0], [0, 14], [14, 16], [0, 15], [15, 17]] + 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]] + if candidate.ndim != 2 or candidate.shape[1] != 2: return canvas + for i in range(len(limbSeq)): + for n in range(len(subset)): + limb = limbSeq[i] + if max(limb) >= subset.shape[1]: continue + index = subset[n][np.array(limb)].astype(int) + if -1 in index or max(index) >= len(candidate): continue + Y, X = candidate[index, 0] * float(W), candidate[index, 1] * float(H) + mX, mY = np.mean(X), np.mean(Y) + length = np.linalg.norm(np.array([X[0], Y[0]]) - np.array([X[1], Y[1]])) + angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) + if length < 1: continue + polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), pose_marker_size), int(angle), 0, 360, 1) + cv2.fillConvexPoly(canvas, polygon, colors[i % len(colors)]) + for n in range(len(subset)): + for i in range(subset.shape[1]): + index = int(subset[n][i]) + if index == -1 or index >= len(candidate): continue + x, y = candidate[index][0:2] + x, y = int(x * W), int(y * H) + cv2.circle(canvas, (x, y), pose_marker_size, colors[i % len(colors)], thickness=-1) + return canvas + +def draw_handpose(canvas, all_hand_peaks, hand_marker_size): + H, W, C = canvas.shape + edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] + for peaks_list_for_one_hand in all_hand_peaks: + peaks_np = np.array(peaks_list_for_one_hand) + if peaks_np.ndim != 2 or peaks_np.shape[1] != 2: continue + for ie, e in enumerate(edges): + if e[0] >= len(peaks_np) or e[1] >= len(peaks_np): continue + x1_coord, y1_coord = peaks_np[e[0]] + x2_coord, y2_coord = peaks_np[e[1]] + if x1_coord < eps and y1_coord < eps or x2_coord < eps and y2_coord < eps: continue + x1, y1 = int(x1_coord * W), int(y1_coord * H) + x2, y2 = int(x2_coord * W), int(y2_coord * H) + if x1 > eps and y1 > eps and x2 > eps and y2 > eps: + cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=max(1, hand_marker_size)) + for i, keyponit in enumerate(peaks_np): + x_coord, y_coord = keyponit + x, y = int(x_coord * W), int(y_coord * H) + if x > eps and y > eps: cv2.circle(canvas, (x, y), max(1, hand_marker_size) + 1, (0, 0, 255), thickness=-1) + return canvas + +def draw_facepose(canvas, all_lmks, face_marker_size): + H, W, C = canvas.shape + lmks_np = np.array(all_lmks) + if lmks_np.ndim != 2 or lmks_np.shape[1] != 2: return canvas + for lmk in lmks_np: + x_coord, y_coord = lmk + x, y = int(x_coord * W), int(y_coord * H) + if x > eps and y > eps: cv2.circle(canvas, (x, y), face_marker_size, (255, 255, 255), thickness=-1) + return canvas