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
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