save code

node_wrappers/pose_keypoint_postprocess.py

@@ -270,6 +270,44 @@ def numpy2torch(np_image: np.ndarray) -> torch.Tensor:
     """ [H, W, C] => [B=1, H, W, C]"""
     return torch.from_numpy(np_image.astype(np.float32) / 255).unsqueeze(0)
 
+class CheckPeopleKps:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "kps": ("POSE_KEYPOINT",)
+            }
+        }
+
+    RETURN_TYPES = ("BOOLEAN", "BOOLEAN")
+    RETURN_NAMES = ("is_leg_ok", "is_hands_ok")
+    FUNCTION = "evaluate"
+    CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
+
+    def evaluate(self, kps):
+        poses, _, height, width = decode_json_as_poses(kps)
+        is_leg_ok = True
+        is_hands_ok = True  
+        import math
+        def distance(p1, p2):
+            return math.sqrt((p1.x - p2.x)**2 + (p1.y - p2.y)**2)
+        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)
+        keypoints = max_pose.body.keypoints
+        unit_len = distance(keypoints[0], keypoints[1])
+        
+        return (is_leg_ok, is_hands_ok)
+
 class RenderPeopleKpsMask:
     @classmethod
     def INPUT_TYPES(s):
@@ -281,7 +319,7 @@ class RenderPeopleKpsMask:
                     "default": "yao"  # 默认值
                 }),
                 "line_len": ("FLOAT", {
-                    "default": 0.5,  # 默认值
+                    "default": 0.4,  # 默认值
                     "min": 0.1,     # 最小值
                     "max": 1.0,      # 最大值
                     "step": 0.01     # 步进值
@@ -399,6 +437,7 @@ NODE_CLASS_MAPPINGS = {
     "RenderPeopleKps": RenderPeopleKps,
     "MyRenderPeopleKps": MyRenderPeopleKps,
     "RenderPeopleKpsMask": RenderPeopleKpsMask,
+    "CheckPeopleKps": CheckPeopleKps,
     "RenderAnimalKps": RenderAnimalKps,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -408,5 +447,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "RenderPeopleKps": "Render Pose JSON (Human)",
     "MyRenderPeopleKps": "My RenderPeopleKps",
     "RenderPeopleKpsMask": "My RenderPeopleKps Mask",
+    "CheckPeopleKps": "My CheckPeopleKps",
     "RenderAnimalKps": "Render Pose JSON (Animal)",
 }

src/custom_controlnet_aux/dwpose/__init__.py

@@ -241,19 +241,20 @@ def distance(p1, p2):
 
 
 def addBase(points, keypoints, H, W, unit_len, closh_short):
-    add_point(points, keypoints, 2, H, W, -unit_len*0.4, -unit_len*0.1)
+    add_point(points, keypoints, 2, H, W, -unit_len*0.4, -unit_len*0.2)
-    add_point(points, keypoints, 1, H, W, 0, -unit_len*0.25)
+    add_point(points, keypoints, 1, H, W, 0, -unit_len*0.4)
-    add_point(points, keypoints, 5, H, W, unit_len*0.4, -unit_len*0.1)
+    add_point(points, keypoints, 5, H, W, unit_len*0.4, -unit_len*0.2)
     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:
+    # if not closh_short:
-        add_point(points, keypoints, 4, H, W, 0, 0)
+    add_point(points, keypoints, 4, H, W, 0, 0)
-        add_point(points, keypoints, 7, 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} 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)
@@ -274,8 +275,8 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
             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, 12, H, W, unit_len*offsetX, 0)
-            add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
+            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)
@@ -284,10 +285,10 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
             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, 12, H, W, unit_len*offsetX, 0)
-            add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
+            add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 0)
-            add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len*1.2)
+            add_point(points, keypoints, 12, H, W, unit_len*offsetX, unit_len*1.2)
-            add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*1.2)
+            add_point(points, keypoints, 9, H, W, -unit_len*offsetX, unit_len*1.2)
     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)
@@ -296,14 +297,14 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
             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, 12, H, W, unit_len*offsetX, 0)
-            add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
+            add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 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, 12, H, W, unit_len*offsetX, unit_len*2.2)
-                add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*2.2)
+                add_point(points, keypoints, 9, H, W, -unit_len*offsetX, unit_len*2.2)
             else:
-                add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*1.2)
+                add_point(points, keypoints, 13, H, W, unit_len*offsetX, unit_len*1.2)
-                add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*1.2)
+                add_point(points, keypoints, 10, H, W, -unit_len*offsetX, unit_len*1.2)
     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)
@@ -312,21 +313,21 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
             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, 12, H, W, unit_len*offsetX, 0)
-            add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
+            add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 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, 12, H, W, unit_len*offsetX, unit_len*3.2)
-                add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*3.2)
+                add_point(points, keypoints, 9, H, W, -unit_len*offsetX, unit_len*3.2)
             else:
-                add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*1)
+                add_point(points, keypoints, 13, H, W, unit_len*offsetX, unit_len*1)
-                add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*1)
+                add_point(points, keypoints, 10, H, W, -unit_len*offsetX, unit_len*1)
     else:
         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)
-    line_len = 0.7
+    # line_len = 0.7
     for i in range(len(points)):
         if i == len(points) - 1:
             p1 = points[i]