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