Add ConvertOpenPoseKeypointsToDWPose node to use dwpose estimator as alternative and for multiple people

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",
 }

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

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: