diff --git a/node_wrappers/pose_keypoint_postprocess.py b/node_wrappers/pose_keypoint_postprocess.py index 83053f9..2bebc38 100644 --- a/node_wrappers/pose_keypoint_postprocess.py +++ b/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)", } diff --git a/src/custom_controlnet_aux/dwpose/__init__.py b/src/custom_controlnet_aux/dwpose/__init__.py index ee90929..e858f87 100644 --- a/src/custom_controlnet_aux/dwpose/__init__.py +++ b/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, 1, H, W, 0, -unit_len*0.25) - add_point(points, keypoints, 5, 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.4) + 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: - add_point(points, keypoints, 4, H, W, 0, 0) - add_point(points, keypoints, 7, H, W, 0, 0) + # 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} 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, 9, 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*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, 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) + add_point(points, keypoints, 12, H, W, unit_len*offsetX, 0) + add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 0) + add_point(points, keypoints, 12, H, W, unit_len*offsetX, 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, 9, 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*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, 9, 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*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, 10, 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*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, 9, 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*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, 9, 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*offsetX, 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) + add_point(points, keypoints, 13, H, W, unit_len*offsetX, 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]