diff --git a/node_wrappers/pose_keypoint_postprocess.py b/node_wrappers/pose_keypoint_postprocess.py index c02db74..55b36db 100644 --- a/node_wrappers/pose_keypoint_postprocess.py +++ b/node_wrappers/pose_keypoint_postprocess.py @@ -327,27 +327,48 @@ class RenderPeopleKpsMask: "cloth_short": ("BOOLEAN", { "default": True }), + "width": ("INT", { + "default": 0 + }), + "height": ("INT", { + "default": 0 + }), + "left_right_line_width": ("FLOAT", { + "default": 0.4 + }), + "human_head2_yao_times": ("FLOAT", { + "default": 2.5 + }), } } - RETURN_TYPES = ("MASK","INT",) + RETURN_TYPES = ("MASK","INT","FLOAT","INT","INT", "INT","INT","INT","INT") + RETURN_NAMES = ("MASK","offset","Head2YaoTimes","Left","Right", "w", "h", "x", "y") FUNCTION = "render" CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess" - def render(self, kps, cloth_len, line_len, cloth_short) -> tuple[torch.Tensor]: + def render(self, kps, cloth_len, line_len, cloth_short, width, height, left_right_line_width, human_head2_yao_times) -> tuple[torch.Tensor]: if isinstance(kps, list): kps = kps[0] poses, _, height, width = decode_json_as_poses(kps) - mask, neck_y = draw_poses_mask( + mask, neck_y, Head2YaoTimes, Left, Right, head_y, head2yao_base = draw_poses_mask( poses, height, width, cloth_len, line_len, cloth_short, + width, + height, + left_right_line_width ) - return (mask, int(neck_y), ) + w = Right - Left + x = Left + h = head2yao_base * human_head2_yao_times + y = head_y + print(f"RenderPeopleKpsMask w:{w} h{h} x{x} y{y}") + return (mask, int(neck_y), float(Head2YaoTimes), int(Left), int(Right), int(w), int(h), int(x), int(y)) class RenderPeopleKps: @classmethod diff --git a/src/custom_controlnet_aux/dwpose/__init__.py b/src/custom_controlnet_aux/dwpose/__init__.py index a76d947..66ca042 100644 --- a/src/custom_controlnet_aux/dwpose/__init__.py +++ b/src/custom_controlnet_aux/dwpose/__init__.py @@ -251,26 +251,42 @@ def addBase(points, keypoints, H, W, unit_len, closh_short): add_point(points, keypoints, 7, H, W, 0, 0) def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, closh_short=True): + for i, kp in enumerate(keypoints): + if kp is None: + print(f" [{i:2d}] ") + else: + # 使用 getattr 安全访问属性,提供默认值 + x = getattr(kp, 'x', 'N/A') + y = getattr(kp, 'y', 'N/A') + score = getattr(kp, 'score', 'N/A') + id_ = getattr(kp, 'id', 'N/A') + + # 格式化输出,处理可能的非数值情况 + x_str = f"{x:8.2f}" if isinstance(x, (int, float)) else f"{x:>8}" + y_str = f"{y:8.2f}" if isinstance(y, (int, float)) else f"{y:>8}" + score_str = f"{score:.2f}" if isinstance(score, (int, float)) else f"{score}" + + print(f" [{i:2d}] x={x_str}, y={y_str}, score={score_str}, id={id_}") points = [] unit_len = distance(keypoints[0], keypoints[1]) print(f'unit_len:{unit_len} cloth_len:{cloth_len}') offsetX = 0.8 if cloth_len == "腰": addBase(points, keypoints, H, W, unit_len, closh_short) - add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*0.3) - add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*0.3) + add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*0.5) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*0.5) elif cloth_len == '跨': addBase(points, keypoints, H, W, unit_len, closh_short) - add_point(points, keypoints, 11, H, W, unit_len*0.2, 0) - add_point(points, keypoints, 8, H, W, -unit_len*0.2, 0) + add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*0.8) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*0.8) elif cloth_len == '大腿': addBase(points, keypoints, H, W, unit_len, closh_short) - add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.5) - add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.5) + add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.8) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.8) elif cloth_len == '膝盖': addBase(points, keypoints, H, W, unit_len, closh_short) - add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.5) - add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.5) + add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.8) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.8) if keypoints[9] == None or keypoints[12] == None: add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*2.2) add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*2.2) @@ -279,8 +295,8 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 0) elif cloth_len == '小腿': addBase(points, keypoints, H, W, unit_len, closh_short) - add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.5) - add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.5) + add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.8) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.8) if keypoints[9] == None or keypoints[12] == None: add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*3) add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*3) @@ -341,7 +357,32 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo return mask -def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_short): +def GetHead2YaoTimes(keypoints, width, height): + yaoY = (keypoints[8].y + keypoints[11].y)/2 + headY = keypoints[0].y + base = yaoY - headY + edge = height - headY + return edge/base, base + +def GetLeftRight(keypoints, width, height, left_right_line_width): + left = 999999 + right = 0 + for i, kp in enumerate(keypoints): + if kp is not None: + if kp.x < left: + left = kp.x + if kp.x > right: + right = kp.x + unit_len = distance(keypoints[1], keypoints[0]) + left = left - unit_len*left_right_line_width + right = right + unit_len*left_right_line_width + if left < 0: + left = 0 + if right >= width: + right = width -1 + return left, right + +def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_short, width, height, left_right_line_width): mask = torch.zeros((H, W), dtype=torch.float32) @@ -359,8 +400,11 @@ def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_sh max_pose = findMaxPose(poses) mask = draw_pose_mask(mask, max_pose.body.keypoints, H, W, cloth_len, line_len, cloth_short) neck_y = max_pose.body.keypoints[1].y - - return mask, neck_y + Head2YaoTimes, head2yao_base = GetHead2YaoTimes(max_pose.body.keypoints, width, height) + Left, Right = GetLeftRight(max_pose.body.keypoints, width, height, left_right_line_width) + head_y = max_pose.body.keypoints[0].y + print(f"Head2YaoTimes:{Head2YaoTimes} Left:{Left} Right:{Right}, head_y:{head_y}, head2yao_base:{head2yao_base}") + return mask, neck_y, Head2YaoTimes, Left, Right, head_y, head2yao_base def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, draw_face=True, xinsr_stick_scaling=False): """