save code

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,)
 

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, 5, H, W, unit_len*0.4, 0)
-        add_point(points, keypoints, 6, H, W)
+        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, 5, H, W, unit_len*0.4, 0)
-        add_point(points, keypoints, 6, H, W)
+        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