save code
This commit is contained in:
@@ -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)",
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user