diff --git a/nodes.py b/nodes.py index cca0d63..fe3af64 100644 --- a/nodes.py +++ b/nodes.py @@ -17,6 +17,68 @@ folder_paths.add_model_folder_path("detection", os.path.join(folder_paths.models from .vitpose_utils.utils import bbox_from_detector, crop, load_pose_metas_from_kp2ds_seq, aaposemeta_to_dwpose_scail +def convert_openpose_to_target_format(frames, max_people=2): + NUM_BODY = 18 + NUM_FACE = 70 + NUM_HAND = 21 + + results = [] + for frame in frames: + canvas_width = frame['canvas_width'] + canvas_height = frame['canvas_height'] + people = frame['people'][:max_people] + + bodies = [] + hands = [] + faces = [] + body_scores = [] + hand_scores = [] + face_scores = [] + + for person in people: + pose_raw = person.get('pose_keypoints_2d') or [] + if len(pose_raw) != NUM_BODY * 3: + continue + + pose = np.array(pose_raw).reshape(-1, 3) + pose_xy = np.stack([pose[:, 0] / canvas_width, pose[:, 1] / canvas_height], axis=1) + bodies.append(pose_xy) + body_scores.append(pose[:, 2]) + + face_raw = person.get('face_keypoints_2d') or [] + if len(face_raw) == NUM_FACE * 3: + face = np.array(face_raw).reshape(-1, 3) + face_xy = np.stack([face[:, 0] / canvas_width, face[:, 1] / canvas_height], axis=1) + faces.append(face_xy) + face_scores.append(face[:, 2]) + + hand_left_raw = person.get('hand_left_keypoints_2d') or [] + hand_right_raw = person.get('hand_right_keypoints_2d') or [] + if len(hand_left_raw) == NUM_HAND * 3: + hand_left = np.array(hand_left_raw).reshape(-1, 3) + hand_left_xy = np.stack([hand_left[:, 0] / canvas_width, hand_left[:, 1] / canvas_height], axis=1) + hands.append(hand_left_xy) + hand_scores.append(hand_left[:, 2]) + if len(hand_right_raw) == NUM_HAND * 3: + hand_right = np.array(hand_right_raw).reshape(-1, 3) + hand_right_xy = np.stack([hand_right[:, 0] / canvas_width, hand_right[:, 1] / canvas_height], axis=1) + hands.append(hand_right_xy) + hand_scores.append(hand_right[:, 2]) + + result = { + 'bodies': { + 'candidate': np.array(bodies, dtype=np.float32), + 'subset': np.array([np.arange(NUM_BODY) for _ in bodies], dtype=np.float32) if bodies else np.array([]) + }, + 'hands': np.array(hands, dtype=np.float32), + 'faces': np.array(faces, dtype=np.float32), + 'body_score': np.array(body_scores, dtype=np.float32), + 'hand_score': np.array(hand_scores, dtype=np.float32), + 'face_score': np.array(face_scores, dtype=np.float32) + } + results.append(result) + return results + def scale_faces(poses, pose_2d_ref): # Input: two lists of dict, poses[0]['faces'].shape: 1, 68, 2 , poses_ref[0]['faces'].shape: 1, 68, 2 # Scale the facial keypoints in poses according to the center point of the face @@ -153,6 +215,27 @@ class PoseDetectionVitPoseToDWPose: return (dwposes,) +class ConvertOpenPoseKeypointsToDWPose: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "keypoints": ("POSE_KEYPOINT",), + "max_people": ("INT", {"default": 2, "min": 1, "max": 100, "step": 1, "tooltip": "Maximum number of people to process per frame"}), + }, + } + + RETURN_TYPES = ("DWPOSES",) + RETURN_NAMES = ("dw_poses",) + FUNCTION = "process" + CATEGORY = "WanAnimatePreprocess" + DESCRIPTION = "Convert OpenPose format keypoints to DWPose format." + + def process(self, keypoints, max_people=2): + + return convert_openpose_to_target_format(keypoints, max_people=max_people), + + class RenderNLFPoses: @classmethod def INPUT_TYPES(s): @@ -202,7 +285,9 @@ class RenderNLFPoses: ori_camera_pose = intrinsic_matrix_from_field_of_view([height, width]) ori_focal = ori_camera_pose[0, 0] - if ref_dw_pose is not None: + num_people = dw_pose_input[0]['bodies']['candidate'].shape[0] if dw_poses is not None else 0 + + if dw_poses is not None and ref_dw_pose is not None and num_people == 1: ref_dw_pose_input = copy.deepcopy(ref_dw_pose) # Find the first valid pose @@ -268,8 +353,10 @@ class RenderNLFPoses: NODE_CLASS_MAPPINGS = { "PoseDetectionVitPoseToDWPose": PoseDetectionVitPoseToDWPose, "RenderNLFPoses": RenderNLFPoses, + "ConvertOpenPoseKeypointsToDWPose": ConvertOpenPoseKeypointsToDWPose, } NODE_DISPLAY_NAME_MAPPINGS = { "PoseDetectionVitPoseToDWPose": "Pose Detection VitPose to DWPose", "RenderNLFPoses": "Render NLF Poses", + "ConvertOpenPoseKeypointsToDWPose": "Convert OpenPose Keypoints to DWPose", } diff --git a/pose_draw/draw_pose_utils.py b/pose_draw/draw_pose_utils.py index 16b4592..753a650 100644 --- a/pose_draw/draw_pose_utils.py +++ b/pose_draw/draw_pose_utils.py @@ -1,7 +1,6 @@ import cv2 import numpy as np from PIL import Image -import os from .draw_utils import draw_bodypose, draw_bodypose_with_feet, draw_handpose_lr, draw_handpose, draw_facepose, draw_bodypose_augmentation @@ -98,16 +97,6 @@ def draw_pose_to_canvas(poses, pool, H, W, reshape_scale, points_only_flag, show return canvas_lst -def get_mp4_filenames_from_directory(dwpose_keypoints_dir): - mp4_filenames_dwpose = [] - # Get all available mp4 files by intersecting keypoints and mp4 - if dwpose_keypoints_dir: - for root, dirs, files in os.walk(dwpose_keypoints_dir): - for file in files: - if file.lower().endswith('.pt'): # Only look for .mp4 files - mp4_filenames_dwpose.append(file.replace(".pt", ".mp4")) # Get absolute path - return mp4_filenames_dwpose - def project_dwpose_to_3d(dwpose_keypoint, original_threed_keypoint, focal, princpt, H, W): # Camera intrinsic parameters # fx, fy = focal, focal diff --git a/pose_draw/draw_utils.py b/pose_draw/draw_utils.py index 243deba..5b30b5a 100644 --- a/pose_draw/draw_utils.py +++ b/pose_draw/draw_utils.py @@ -203,12 +203,12 @@ def draw_bodypose_augmentation(canvas, candidate, subset, drop_aug=True, shift_a stickwidth = 4 limbSeq = [ - [2, 3], # 1->2 左肩 0 - [2, 6], # 1->5 右肩 1 - [3, 4], # 2->3 左臂 2 - [4, 5], # 3->4 左肘 3 - [6, 7], # 5->6 右臂 4 - [7, 8], # 6->7 右肘 5 + [2, 3], # 1->2 left shoulder 0 + [2, 6], # 1->5 right shoulder 1 + [3, 4], # 2->3 left arm 2 + [4, 5], # 3->4 left elbow 3 + [6, 7], # 5->6 right arm 4 + [7, 8], # 6->7 right elbow 5 [2, 9], # 6 [9, 10], # 7 [10, 11], # 8 @@ -245,7 +245,7 @@ def draw_bodypose_augmentation(canvas, candidate, subset, drop_aug=True, shift_a [255, 0, 85], ] - # 随机选0-2根骨骼进行丢弃 + # Randomly select 0-2 bones to drop if drop_aug: arr_drop = list(range(17)) k_drop = random.choices([0, 1, 2], weights=[0.5, 0.3, 0.2])[0] @@ -257,7 +257,7 @@ def draw_bodypose_augmentation(canvas, candidate, subset, drop_aug=True, shift_a else: shift_indices = [] if all_cheek_aug: - drop_indices = list(range(13)) # 0-12对应的骨骼都扔掉 + drop_indices = list(range(13)) # Drop all bones corresponding to 0-12 for i in range(17): for n in range(len(subset)): @@ -270,7 +270,7 @@ def draw_bodypose_augmentation(canvas, candidate, subset, drop_aug=True, shift_a if i in drop_indices: continue - mX = np.mean(X) # 计算两个关节点之间的中点 + mX = np.mean(X) # Calculate the midpoint between two joints mY = np.mean(Y) length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 if i in shift_indices: