save code

src/custom_controlnet_aux/dwpose/__init__.py

@@ -344,8 +344,19 @@ def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_sh
 
     mask = torch.zeros((H, W), dtype=torch.float32)
 
-    for pose in poses:
-        mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len, line_len, cloth_short)
+    def findMaxPose(poses):
+        len = 0
+        mp = None
+        for p in poses:
+            if p.body.keypoints[0] != None and p.body.keypoints[1] != None:
+                l = distance(p.body.keypoints[0], p.body.keypoints[1])
+                if l > len:
+                    mp = p
+                    len = l
+        return mp
+            
+    max_pose = findMaxPose(poses)
+    mask = draw_pose_mask(mask, max_pose.body.keypoints, H, W, cloth_len, line_len, cloth_short)
             
     return mask
 
@@ -366,16 +377,27 @@ def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, dr
     """
     canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
 
-    for pose in poses:
-        if draw_body:
-            canvas = util.draw_bodypose(canvas, pose.body.keypoints, xinsr_stick_scaling)
+    def findMaxPose(poses):
+        len = 0
+        mp = None
+        for p in poses:
+            if p.body.keypoints[0] != None and p.body.keypoints[1] != None:
+                l = distance(p.body.keypoints[0], p.body.keypoints[1])
+                if l > len:
+                    mp = p
+                    len = l
+        return mp
+    pose = findMaxPose(poses)
+    # for pose in poses:
+    if draw_body:
+        canvas = util.draw_bodypose(canvas, pose.body.keypoints, xinsr_stick_scaling)
 
-        if draw_hand:
-            canvas = util.draw_handpose(canvas, pose.left_hand)
-            canvas = util.draw_handpose(canvas, pose.right_hand)
+    if draw_hand:
+        canvas = util.draw_handpose(canvas, pose.left_hand)
+        canvas = util.draw_handpose(canvas, pose.right_hand)
 
-        if draw_face:
-            canvas = util.draw_facepose(canvas, pose.face)
+    if draw_face:
+        canvas = util.draw_facepose(canvas, pose.face)
 
     return canvas