xsl/comfyui_controlnet_aux
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

save code

Browse Source
This commit is contained in:
colomi
2025-07-13 00:25:38 +08:00
parent b0b264b421
commit 0a6a9ecfb3
2 changed files with 68 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node_wrappers/pose_keypoint_postprocess.py
+8 -1
View File
@@ -280,6 +280,12 @@ class RenderPeopleKpsMask:
"multiline": False, # 是否允许多行输入
"default": "yao" # 默认值
}),
"line_len": ("FLOAT", {
"default": 0.5, # 默认值
"min": 0.1, # 最小值
"max": 1.0, # 最大值
"step": 0.01 # 步进值
}),
}
}
@@ -287,7 +293,7 @@ class RenderPeopleKpsMask:
FUNCTION = "render"
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
def render(self, kps, cloth_len) -> tuple[torch.Tensor]:
def render(self, kps, cloth_len, line_len) -> tuple[torch.Tensor]:
if isinstance(kps, list):
kps = kps[0]
@@ -297,6 +303,7 @@ class RenderPeopleKpsMask:
height,
width,
cloth_len,
line_len,
)
return (mask,)
src/custom_controlnet_aux/dwpose/__init__.py
+60 -12
View File
@@ -240,26 +240,74 @@ def distance(p1, p2):
return math.sqrt((p1.x - p2.x)**2 + (p1.y - p2.y)**2)
def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len):
def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len):
points = []
unit_len = distance(keypoints[0], keypoints[1])
print(f'unit_len {unit_len}')
if cloth_len == "腰":
add_point(points, keypoints, 2, H, W, -unit_len*0.2, 0)
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 6, H, W)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
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, 3, H, W)
add_point(points, keypoints, 3, H, W, -unit_len*0.2)
elif cloth_len == '跨':
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
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, 3, H, W, -unit_len*0.2)
elif cloth_len == '大腿':
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len)
add_point(points, keypoints, 3, H, W, -unit_len*0.2)
elif cloth_len == '膝盖':
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
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, 3, H, W, -unit_len*0.2)
elif cloth_len == '小腿':
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len)
add_point(points, keypoints, 3, H, W, -unit_len*0.2)
elif cloth_len == '脚踝':
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
add_point(points, keypoints, 13, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 10, H, W, -unit_len*0.2, 0)
add_point(points, keypoints, 3, H, W, -unit_len*0.2)
elif cloth_len == '拖地':
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*0.5)
add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*0.5)
add_point(points, keypoints, 3, H, W, -unit_len*0.2)
else:
add_point(points, keypoints, 2, H, W, -unit_len*0.2, 0)
add_point(points, keypoints, 2, H, W, -unit_len*0.4, 0)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
add_point(points, keypoints, 5, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 6, H, W)
add_point(points, keypoints, 5, H, W, unit_len*0.4, 0)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
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, 3, H, W)
add_point(points, keypoints, 3, H, W, -unit_len*0.2)
points = convex_hull(points)
@@ -272,16 +320,16 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len):
p2 = points[i + 1]
if p1 is not None and p2 is not None:
draw_pose_mask_line(mask, H, W, p1, p2, line_width=(unit_len*0.1))
draw_pose_mask_line(mask, H, W, p1, p2, line_width=(unit_len*line_len))
return mask
def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len):
def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len):
mask = torch.zeros((H, W), dtype=torch.float32)
for pose in poses:
mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len)
mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len, line_len)
return mask