save code

This commit is contained in:
xsl
2025-08-17 23:35:31 +08:00
parent 88ea7bd6ba
commit e1a03d8e20
2 changed files with 82 additions and 17 deletions
+25 -4
View File
@@ -327,27 +327,48 @@ class RenderPeopleKpsMask:
"cloth_short": ("BOOLEAN", {
"default": True
}),
"width": ("INT", {
"default": 0
}),
"height": ("INT", {
"default": 0
}),
"left_right_line_width": ("FLOAT", {
"default": 0.4
}),
"human_head2_yao_times": ("FLOAT", {
"default": 2.5
}),
}
}
RETURN_TYPES = ("MASK","INT",)
RETURN_TYPES = ("MASK","INT","FLOAT","INT","INT", "INT","INT","INT","INT")
RETURN_NAMES = ("MASK","offset","Head2YaoTimes","Left","Right", "w", "h", "x", "y")
FUNCTION = "render"
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
def render(self, kps, cloth_len, line_len, cloth_short) -> tuple[torch.Tensor]:
def render(self, kps, cloth_len, line_len, cloth_short, width, height, left_right_line_width, human_head2_yao_times) -> tuple[torch.Tensor]:
if isinstance(kps, list):
kps = kps[0]
poses, _, height, width = decode_json_as_poses(kps)
mask, neck_y = draw_poses_mask(
mask, neck_y, Head2YaoTimes, Left, Right, head_y, head2yao_base = draw_poses_mask(
poses,
height,
width,
cloth_len,
line_len,
cloth_short,
width,
height,
left_right_line_width
)
return (mask, int(neck_y), )
w = Right - Left
x = Left
h = head2yao_base * human_head2_yao_times
y = head_y
print(f"RenderPeopleKpsMask w:{w} h{h} x{x} y{y}")
return (mask, int(neck_y), float(Head2YaoTimes), int(Left), int(Right), int(w), int(h), int(x), int(y))
class RenderPeopleKps:
@classmethod
+57 -13
View File
@@ -251,26 +251,42 @@ def addBase(points, keypoints, H, W, unit_len, closh_short):
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):
for i, kp in enumerate(keypoints):
if kp is None:
print(f" [{i:2d}] <Missing Keypoint>")
else:
# 使用 getattr 安全访问属性,提供默认值
x = getattr(kp, 'x', 'N/A')
y = getattr(kp, 'y', 'N/A')
score = getattr(kp, 'score', 'N/A')
id_ = getattr(kp, 'id', 'N/A')
# 格式化输出,处理可能的非数值情况
x_str = f"{x:8.2f}" if isinstance(x, (int, float)) else f"{x:>8}"
y_str = f"{y:8.2f}" if isinstance(y, (int, float)) else f"{y:>8}"
score_str = f"{score:.2f}" if isinstance(score, (int, float)) else f"{score}"
print(f" [{i:2d}] x={x_str}, y={y_str}, score={score_str}, id={id_}")
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)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*0.3)
add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*0.5)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*0.5)
elif cloth_len == '':
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)
add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*0.8)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*0.8)
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)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.5)
add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.8)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.8)
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)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.5)
add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.8)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.8)
if keypoints[9] == None or keypoints[12] == None:
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)
@@ -279,8 +295,8 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
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)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.5)
add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.8)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*1.8)
if keypoints[9] == None or keypoints[12] == None:
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)
@@ -341,7 +357,32 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
return mask
def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_short):
def GetHead2YaoTimes(keypoints, width, height):
yaoY = (keypoints[8].y + keypoints[11].y)/2
headY = keypoints[0].y
base = yaoY - headY
edge = height - headY
return edge/base, base
def GetLeftRight(keypoints, width, height, left_right_line_width):
left = 999999
right = 0
for i, kp in enumerate(keypoints):
if kp is not None:
if kp.x < left:
left = kp.x
if kp.x > right:
right = kp.x
unit_len = distance(keypoints[1], keypoints[0])
left = left - unit_len*left_right_line_width
right = right + unit_len*left_right_line_width
if left < 0:
left = 0
if right >= width:
right = width -1
return left, right
def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_short, width, height, left_right_line_width):
mask = torch.zeros((H, W), dtype=torch.float32)
@@ -359,8 +400,11 @@ def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_sh
max_pose = findMaxPose(poses)
mask = draw_pose_mask(mask, max_pose.body.keypoints, H, W, cloth_len, line_len, cloth_short)
neck_y = max_pose.body.keypoints[1].y
return mask, neck_y
Head2YaoTimes, head2yao_base = GetHead2YaoTimes(max_pose.body.keypoints, width, height)
Left, Right = GetLeftRight(max_pose.body.keypoints, width, height, left_right_line_width)
head_y = max_pose.body.keypoints[0].y
print(f"Head2YaoTimes:{Head2YaoTimes} Left:{Left} Right:{Right}, head_y:{head_y}, head2yao_base:{head2yao_base}")
return mask, neck_y, Head2YaoTimes, Left, Right, head_y, head2yao_base
def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, draw_face=True, xinsr_stick_scaling=False):
"""