From 0a6a9ecfb3f7e80c374a4ac987370bf0dd53b594 Mon Sep 17 00:00:00 2001 From: colomi <1421901449@qq.com> Date: Sun, 13 Jul 2025 00:25:38 +0800 Subject: [PATCH] save code --- node_wrappers/pose_keypoint_postprocess.py | 9 ++- src/custom_controlnet_aux/dwpose/__init__.py | 72 ++++++++++++++++---- 2 files changed, 68 insertions(+), 13 deletions(-) diff --git a/node_wrappers/pose_keypoint_postprocess.py b/node_wrappers/pose_keypoint_postprocess.py index 7bc6184..5be2c29 100644 --- a/node_wrappers/pose_keypoint_postprocess.py +++ b/node_wrappers/pose_keypoint_postprocess.py @@ -280,6 +280,12 @@ class RenderPeopleKpsMask: "multiline": False, # 是否允许多行输入 "default": "yao" # 默认值 }), + "line_len": ("FLOAT", { + "default": 0.5, # 默认值 + "min": 0.1, # 最小值 + "max": 1.0, # 最大值 + "step": 0.01 # 步进值 + }), } } @@ -287,7 +293,7 @@ class RenderPeopleKpsMask: FUNCTION = "render" CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess" - def render(self, kps, cloth_len) -> tuple[torch.Tensor]: + def render(self, kps, cloth_len, line_len) -> tuple[torch.Tensor]: if isinstance(kps, list): kps = kps[0] @@ -297,6 +303,7 @@ class RenderPeopleKpsMask: height, width, cloth_len, + line_len, ) return (mask,) diff --git a/src/custom_controlnet_aux/dwpose/__init__.py b/src/custom_controlnet_aux/dwpose/__init__.py index e16bc71..e97ad5b 100644 --- a/src/custom_controlnet_aux/dwpose/__init__.py +++ b/src/custom_controlnet_aux/dwpose/__init__.py @@ -240,26 +240,74 @@ def distance(p1, p2): return math.sqrt((p1.x - p2.x)**2 + (p1.y - p2.y)**2) -def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len): +def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len): points = [] unit_len = distance(keypoints[0], keypoints[1]) print(f'unit_len {unit_len}') if cloth_len == "腰": - add_point(points, keypoints, 2, H, W, -unit_len*0.2, 0) + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) - add_point(points, keypoints, 5, H, W, unit_len*0.2, 0) - add_point(points, keypoints, 6, H, W) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) 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, 3, H, W) + add_point(points, keypoints, 3, H, W, -unit_len*0.2) + elif cloth_len == '跨': + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) + add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) + 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, 3, H, W, -unit_len*0.2) + elif cloth_len == '大腿': + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) + add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) + add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len) + add_point(points, keypoints, 3, H, W, -unit_len*0.2) + elif cloth_len == '膝盖': + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) + add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) + add_point(points, keypoints, 12, H, W, unit_len*0.2, 0) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0) + add_point(points, keypoints, 3, H, W, -unit_len*0.2) + elif cloth_len == '小腿': + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) + add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) + add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len) + add_point(points, keypoints, 3, H, W, -unit_len*0.2) + elif cloth_len == '脚踝': + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) + add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) + add_point(points, keypoints, 13, H, W, unit_len*0.2, 0) + add_point(points, keypoints, 10, H, W, -unit_len*0.2, 0) + add_point(points, keypoints, 3, H, W, -unit_len*0.2) + elif cloth_len == '拖地': + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) + add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) + add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*0.5) + add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*0.5) + add_point(points, keypoints, 3, H, W, -unit_len*0.2) else: - add_point(points, keypoints, 2, H, W, -unit_len*0.2, 0) + add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0) add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2) - add_point(points, keypoints, 5, H, W, unit_len*0.2, 0) - add_point(points, keypoints, 6, H, W) + add_point(points, keypoints, 5, H, W, unit_len*0.4, 0) + add_point(points, keypoints, 6, H, W, unit_len*0.2) 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, 3, H, W) + add_point(points, keypoints, 3, H, W, -unit_len*0.2) points = convex_hull(points) @@ -272,16 +320,16 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len): p2 = points[i + 1] if p1 is not None and p2 is not None: - draw_pose_mask_line(mask, H, W, p1, p2, line_width=(unit_len*0.1)) + draw_pose_mask_line(mask, H, W, p1, p2, line_width=(unit_len*line_len)) return mask -def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len): +def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len): mask = torch.zeros((H, W), dtype=torch.float32) for pose in poses: - mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len) + mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len, line_len) return mask