nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Make dwpose keypoint size adjustable

Browse Source
This commit is contained in:
kijai
2025-04-19 20:08:30 +03:00
parent 7b6c34e26d
commit b5b4e44512
2 changed files with 21 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
unianimate/dwpose/util.py
+6 -6
View File
@@ -109,7 +109,7 @@ def draw_bodypose(canvas, candidate, subset):
return canvas
def draw_body_and_foot(canvas, candidate, subset, stick_width=4, draw_body=True, draw_feet=True, draw_body_keypoints=True):
def draw_body_and_foot(canvas, candidate, subset, stick_width=4, draw_body=True, draw_feet=True, body_keypoint_size=4):
H, W, C = canvas.shape
candidate = np.array(candidate)
subset = np.array(subset)
@@ -144,7 +144,7 @@ def draw_body_and_foot(canvas, candidate, subset, stick_width=4, draw_body=True,
canvas = (canvas * 0.6).astype(np.uint8)
if draw_body_keypoints:
if body_keypoint_size > 0:
for i in range(len(limbSeq)+1):
for n in range(len(subset)):
index = int(subset[n][i])
@@ -153,12 +153,12 @@ def draw_body_and_foot(canvas, candidate, subset, stick_width=4, draw_body=True,
x, y = candidate[index][0:2]
x = int(x * W)
y = int(y * H)
cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
cv2.circle(canvas, (int(x), int(y)), body_keypoint_size, colors[i], thickness=-1)
return canvas
def draw_handpose(canvas, all_hand_peaks, draw_hands=True, draw_hand_keypoints=True):
def draw_handpose(canvas, all_hand_peaks, draw_hands=True, hand_keypoint_size=4):
H, W, C = canvas.shape
edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
@@ -178,13 +178,13 @@ def draw_handpose(canvas, all_hand_peaks, draw_hands=True, draw_hand_keypoints=T
if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=2)
if draw_hand_keypoints:
if hand_keypoint_size > 0:
for i, keypoint in enumerate(peaks):
x, y = keypoint
x = int(x * W)
y = int(y * H)
if x > eps and y > eps:
cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
cv2.circle(canvas, (x, y), hand_keypoint_size, (0, 0, 255), thickness=-1)
return canvas
unianimate/nodes.py
+15 -14
View File
@@ -185,7 +185,7 @@ class DWposeDetector:
return pose
def draw_pose(pose, H, W, stick_width=4,draw_body=True, draw_hands=True, draw_feet=True,
draw_body_keypoints=True, draw_hand_keypoints=True):
body_keypoint_size=4, hand_keypoint_size=4):
bodies = pose['bodies']
faces = pose['faces']
hands = pose['hands']
@@ -193,8 +193,8 @@ def draw_pose(pose, H, W, stick_width=4,draw_body=True, draw_hands=True, draw_fe
subset = bodies['subset']
canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
canvas = draw_body_and_foot(canvas, candidate, subset, draw_body=draw_body, stick_width=stick_width, draw_feet=draw_feet, draw_body_keypoints=draw_body_keypoints)
canvas = draw_handpose(canvas, hands, draw_hands=draw_hands, draw_hand_keypoints=draw_hand_keypoints)
canvas = draw_body_and_foot(canvas, candidate, subset, draw_body=draw_body, stick_width=stick_width, draw_feet=draw_feet, body_keypoint_size=body_keypoint_size)
canvas = draw_handpose(canvas, hands, draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size)
canvas_without_face = copy.deepcopy(canvas)
canvas = draw_facepose(canvas, faces)
@@ -202,8 +202,8 @@ def draw_pose(pose, H, W, stick_width=4,draw_body=True, draw_hands=True, draw_fe
def pose_extract(pose_images, ref_image, dwpose_model, height, width, score_threshold, stick_width,
draw_body=True, draw_hands=True, draw_hand_keypoints=True, draw_feet=True,
draw_body_keypoints=True):
draw_body=True, draw_hands=True, hand_keypoint_size=4, draw_feet=True,
body_keypoint_size=4):
results_vis = []
comfy_pbar = ProgressBar(len(pose_images))
@@ -678,8 +678,8 @@ def pose_extract(pose_images, ref_image, dwpose_model, height, width, score_thre
for i in range(len(results_vis)):
try:
dwpose_woface, dwpose_wface = draw_pose(results_vis[i], H=height, W=width, stick_width=stick_width,
draw_body=draw_body, draw_hands=draw_hands, draw_hand_keypoints=draw_hand_keypoints,
draw_feet=draw_feet, draw_body_keypoints=draw_body_keypoints)
draw_body=draw_body, draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size,
draw_feet=draw_feet, body_keypoint_size=body_keypoint_size)
result = torch.from_numpy(dwpose_woface)
except:
result = torch.zeros((height, width, 3), dtype=torch.uint8)
@@ -689,8 +689,8 @@ def pose_extract(pose_images, ref_image, dwpose_model, height, width, score_thre
dwpose_woface_ref_tensor = None
if ref_image is not None:
dwpose_woface_ref, dwpose_wface_ref = draw_pose(pose_ref, H=height, W=width, stick_width=stick_width,
draw_body=draw_body, draw_hands=draw_hands, draw_hand_keypoints=draw_hand_keypoints,
draw_feet=draw_feet, draw_body_keypoints=draw_body_keypoints)
draw_body=draw_body, draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size,
draw_feet=draw_feet, body_keypoint_size=body_keypoint_size)
dwpose_woface_ref_tensor = torch.from_numpy(dwpose_woface_ref)
return dwpose_woface_tensor, dwpose_woface_ref_tensor
@@ -703,10 +703,10 @@ class WanVideoUniAnimateDWPoseDetector:
"score_threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Score threshold for pose detection"}),
"stick_width": ("INT", {"default": 4, "min": 1, "max": 100, "step": 1, "tooltip": "Stick width for drawing keypoints"}),
"draw_body": ("BOOLEAN", {"default": True, "tooltip": "Draw body keypoints"}),
"draw_body_keypoints": ("BOOLEAN", {"default": True, "tooltip": "Draw body keypoints"}),
"body_keypoint_size": ("INT", {"default": 4, "min": 0, "max": 100, "step": 1, "tooltip": "Body keypoint size"}),
"draw_feet": ("BOOLEAN", {"default": True, "tooltip": "Draw feet keypoints"}),
"draw_hands": ("BOOLEAN", {"default": True, "tooltip": "Draw hand keypoints"}),
"draw_hand_keypoints": ("BOOLEAN", {"default": True, "tooltip": "Draw hand keypoints"}),
"hand_keypoint_size": ("INT", {"default": 4, "min": 0, "max": 100, "step": 1, "tooltip": "Hand keypoint size"}),
},
"optional": {
"reference_pose_image": ("IMAGE", {"tooltip": "Reference pose image"}),
@@ -718,7 +718,8 @@ class WanVideoUniAnimateDWPoseDetector:
FUNCTION = "process"
CATEGORY = "WanVideoWrapper"
def process(self, pose_images, score_threshold, stick_width, reference_pose_image=None, draw_body=True, draw_body_keypoints=True, draw_feet=True, draw_hands=True, draw_hand_keypoints=True):
def process(self, pose_images, score_threshold, stick_width, reference_pose_image=None, draw_body=True, body_keypoint_size=4,
draw_feet=True, draw_hands=True, hand_keypoint_size=4):
device = mm.get_torch_device()
@@ -763,8 +764,8 @@ class WanVideoUniAnimateDWPoseDetector:
ref_np = ref.cpu().numpy() * 255
poses, reference_pose = pose_extract(pose_np, ref_np, self.dwpose_detector, height, width, score_threshold, stick_width=stick_width,
draw_body=draw_body, draw_body_keypoints=draw_body_keypoints, draw_feet=draw_feet,
draw_hands=draw_hands, draw_hand_keypoints=draw_hand_keypoints)
draw_body=draw_body, body_keypoint_size=body_keypoint_size, draw_feet=draw_feet,
draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size)
poses = poses / 255.0
if reference_pose_image is not None:
reference_pose = reference_pose.unsqueeze(0) / 255.0