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 23:55:04 +08:00
parent 0a6a9ecfb3
commit 312f40d5e0
2 changed files with 77 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node_wrappers/pose_keypoint_postprocess.py
+5 -1
View File
@@ -286,6 +286,9 @@ class RenderPeopleKpsMask:
"max": 1.0, # 最大值
"step": 0.01 # 步进值
}),
"cloth_short": ("BOOLEAN", {
"default": True
}),
}
}
@@ -293,7 +296,7 @@ class RenderPeopleKpsMask:
FUNCTION = "render"
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
def render(self, kps, cloth_len, line_len) -> tuple[torch.Tensor]:
def render(self, kps, cloth_len, line_len, cloth_short) -> tuple[torch.Tensor]:
if isinstance(kps, list):
kps = kps[0]
@@ -304,6 +307,7 @@ class RenderPeopleKpsMask:
width,
cloth_len,
line_len,
cloth_short,
)
return (mask,)
src/custom_controlnet_aux/dwpose/__init__.py
+72 -56
View File
@@ -240,74 +240,90 @@ 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, line_len):
def addBase(points, keypoints, H, W, unit_len, closh_short):
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, 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)
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}')
print(f'unit_len:{unit_len} cloth_len:{cloth_len}')
if 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)
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, 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)
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, 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)
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)
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)
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)
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)
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, 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)
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)
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)
else:
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*0.3)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*0.3)
add_point(points, keypoints, 3, 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, 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)
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)
if keypoints[9] == None or keypoints[12] == None:
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)
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)
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)
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)
if keypoints[9] == None or keypoints[12] == None:
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)
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)
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)
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)
@@ -324,12 +340,12 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len):
return mask
def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len):
def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len, line_len, cloth_short):
mask = torch.zeros((H, W), dtype=torch.float32)
for pose in poses:
mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len, line_len)
mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len, line_len, cloth_short)
return mask