diff --git a/src/custom_controlnet_aux/dwpose/__init__.py b/src/custom_controlnet_aux/dwpose/__init__.py index bb3ac06..ee90929 100644 --- a/src/custom_controlnet_aux/dwpose/__init__.py +++ b/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 diff --git a/src/custom_controlnet_aux/dwpose/__pycache__/__init__.cpython-312.pyc b/src/custom_controlnet_aux/dwpose/__pycache__/__init__.cpython-312.pyc index 3afe575..c6db69a 100644 Binary files a/src/custom_controlnet_aux/dwpose/__pycache__/__init__.cpython-312.pyc and b/src/custom_controlnet_aux/dwpose/__pycache__/__init__.cpython-312.pyc differ