wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-23 16:29:25 +08:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

Draw head without body

Browse Source
This commit is contained in:
rma 2026-06-14 13:33:50 +03:00
parent 1eaeeb75d4
commit d3529a38c7
1 changed files with 25 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
comfy_extras/nodes_sdpose.py
View File

@ -96,8 +96,12 @@ class KeypointDraw:
# Body connections - matching DWPose limbSeq (1-indexed, converted to 0-indexed) # Body connections - matching DWPose limbSeq (1-indexed, converted to 0-indexed)
self.body_limbSeq = [ self.body_limbSeq = [
[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], [2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10],
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], [10, 11], [2, 12], [12, 13], [13, 14]
[1, 16], [16, 18] ]
# Head connections (1-indexed, converted to 0-indexed)
self.head_edges = [
[2, 1], [1, 15], [15, 17], [1, 16], [16, 18]
] ]
# Colors matching DWPose # Colors matching DWPose
@ -237,12 +241,17 @@ class KeypointDraw:
""" """
H, W, C = canvas.shape H, W, C = canvas.shape
# Draw body limbs # Draw body limbs & head connections
if draw_body and len(keypoints) >= 18: if (draw_body or draw_head) and len(keypoints) >= 18:
for i, limb in enumerate(self.body_limbSeq): colorIndexOffset = 0
# Skip head limbs edges = []
if not draw_head and i > 11: if draw_body:
continue edges += self.body_limbSeq
else:
colorIndexOffset += len(self.body_limbSeq)
if draw_head:
edges += self.head_edges
for i, limb in enumerate(edges):
# Convert from 1-indexed to 0-indexed # Convert from 1-indexed to 0-indexed
idx1, idx2 = limb[0] - 1, limb[1] - 1 idx1, idx2 = limb[0] - 1, limb[1] - 1
@ -265,13 +274,16 @@ class KeypointDraw:
polygon = self.draw.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stick_width), int(angle), 0, 360, 1) polygon = self.draw.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stick_width), int(angle), 0, 360, 1)
self.draw.fillConvexPoly(canvas, polygon, self.colors[i % len(self.colors)]) self.draw.fillConvexPoly(canvas, polygon, self.colors[(i + colorIndexOffset) % len(self.colors)])
# Draw body keypoints # Draw body & head keypoints
if draw_body and len(keypoints) >= 18: if (draw_body or draw_head) and len(keypoints) >= 18:
head_keypoints = {0, 14, 15, 16, 17} # nose, eyes, ears
neck_point = 1
for i in range(18): for i in range(18):
# Skip head points (0 - nose) if not draw_head and i in head_keypoints:
if not draw_head and (i == 0 or i > 13): continue
if not draw_body and i not in head_keypoints and i != neck_point:
continue continue
if scores is not None and scores[i] < threshold: if scores is not None and scores[i] < threshold:
continue continue