From b5b4e44512623851b1de75cac7e972b57058ef29 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sat, 19 Apr 2025 20:08:30 +0300 Subject: [PATCH] Make dwpose keypoint size adjustable --- unianimate/dwpose/util.py | 12 ++++++------ unianimate/nodes.py | 29 +++++++++++++++-------------- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/unianimate/dwpose/util.py b/unianimate/dwpose/util.py index bc8a6e3..f809094 100644 --- a/unianimate/dwpose/util.py +++ b/unianimate/dwpose/util.py @@ -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 diff --git a/unianimate/nodes.py b/unianimate/nodes.py index e12129e..c02c7b1 100644 --- a/unianimate/nodes.py +++ b/unianimate/nodes.py @@ -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