From 312f40d5e00b0f4605866d7863bacbe404b5a206 Mon Sep 17 00:00:00 2001 From: colomi <1421901449@qq.com> Date: Sun, 13 Jul 2025 23:55:04 +0800 Subject: [PATCH] save code --- node_wrappers/pose_keypoint_postprocess.py | 6 +- src/custom_controlnet_aux/dwpose/__init__.py | 128 +++++++++++-------- 2 files changed, 77 insertions(+), 57 deletions(-) diff --git a/node_wrappers/pose_keypoint_postprocess.py b/node_wrappers/pose_keypoint_postprocess.py index 5be2c29..83053f9 100644 --- a/node_wrappers/pose_keypoint_postprocess.py +++ b/node_wrappers/pose_keypoint_postprocess.py @@ -286,6 +286,9 @@ class RenderPeopleKpsMask: "max": 1.0, # 最大值 "step": 0.01 # 步进值 }), + "cloth_short": ("BOOLEAN", { + "default": True + }), } } @@ -293,7 +296,7 @@ class RenderPeopleKpsMask: FUNCTION = "render" CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess" - def render(self, kps, cloth_len, line_len) -> tuple[torch.Tensor]: + def render(self, kps, cloth_len, line_len, cloth_short) -> tuple[torch.Tensor]: if isinstance(kps, list): kps = kps[0] @@ -304,6 +307,7 @@ class RenderPeopleKpsMask: width, cloth_len, line_len, + cloth_short, ) return (mask,) diff --git a/src/custom_controlnet_aux/dwpose/__init__.py b/src/custom_controlnet_aux/dwpose/__init__.py index e97ad5b..ee7c6dc 100644 --- a/src/custom_controlnet_aux/dwpose/__init__.py +++ b/src/custom_controlnet_aux/dwpose/__init__.py @@ -240,74 +240,90 @@ 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, line_len): +def addBase(points, keypoints, H, W, unit_len, closh_short): + 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, 3, H, W, -unit_len*0.2) + if not closh_short: + add_point(points, keypoints, 4, H, W, 0, 0) + 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): points = [] unit_len = distance(keypoints[0], keypoints[1]) - print(f'unit_len {unit_len}') + print(f'unit_len:{unit_len} cloth_len:{cloth_len}') if 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) + 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, 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) + 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, 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) + 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) 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) + 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) + 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) + else: + add_point(points, keypoints, 12, H, W, unit_len*0.2, 0) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0) 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) + 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) + 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) + else: + 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, 12, H, W, unit_len*0.2, unit_len*1.2) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*1.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) + 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) + if keypoints[9] == None or keypoints[12] == None: + add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*6) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*6) + else: + add_point(points, keypoints, 12, H, W, unit_len*0.2, 0) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0) + if keypoints[10] == None or keypoints[13] == None: + add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len*2.2) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*2.2) + else: + add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*1.2) + add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*1.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) + 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) + if keypoints[9] == None or keypoints[12] == None: + add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*7) + add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*7) + else: + add_point(points, keypoints, 12, H, W, unit_len*0.2, 0) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0) + if keypoints[10] == None or keypoints[13] == None: + add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len*3.2) + add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*3.2) + else: + add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*1) + add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*1) else: - 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*0.3) - add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*0.3) - add_point(points, keypoints, 3, H, W, -unit_len*0.2) + 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) points = convex_hull(points) @@ -324,12 +340,12 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len): return mask -def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len): +def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_short): 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) + mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len, line_len, cloth_short) return mask