save code

This commit is contained in:
xsl
2025-08-09 14:50:05 +08:00
parent 62ef148f72
commit 2f1e03ca15
2 changed files with 67 additions and 26 deletions
+41 -1
View File
@@ -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)",
}
+26 -25
View File
@@ -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]