From 7f43e61535480ceb739749a90af91b46f7811a71 Mon Sep 17 00:00:00 2001 From: mymusise Date: Wed, 17 Jul 2024 17:36:24 +0800 Subject: [PATCH 1/2] add draw_body, draw_hand, draw_face args Signed-off-by: mymusise --- src/controlnet_aux/dwpose/__init__.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/controlnet_aux/dwpose/__init__.py b/src/controlnet_aux/dwpose/__init__.py index 34e010f..f009223 100644 --- a/src/controlnet_aux/dwpose/__init__.py +++ b/src/controlnet_aux/dwpose/__init__.py @@ -16,7 +16,7 @@ from . import util -def draw_pose(pose, H, W): +def draw_pose(pose, H, W, draw_body=True, draw_hand=True, draw_face=True): bodies = pose['bodies'] faces = pose['faces'] hands = pose['hands'] @@ -24,9 +24,12 @@ def draw_pose(pose, H, W): subset = bodies['subset'] canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8) - canvas = util.draw_bodypose(canvas, candidate, subset) - canvas = util.draw_handpose(canvas, hands) - canvas = util.draw_facepose(canvas, faces) + if draw_body: + canvas = util.draw_bodypose(canvas, candidate, subset) + if draw_hand: + canvas = util.draw_handpose(canvas, hands) + if draw_face: + canvas = util.draw_facepose(canvas, faces) return canvas @@ -40,7 +43,7 @@ def to(self, device): self.pose_estimation.to(device) return self - def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs): + def __call__(self, input_image, detect_resolution=512, image_resolution=512, include_body=True, include_hand=False, include_face=False, output_type="pil", **kwargs): input_image = cv2.cvtColor(np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR) @@ -77,7 +80,7 @@ def __call__(self, input_image, detect_resolution=512, image_resolution=512, out bodies = dict(candidate=body, subset=score) pose = dict(bodies=bodies, hands=hands, faces=faces) - detected_map = draw_pose(pose, H, W) + detected_map = draw_pose(pose, H, W, draw_body=include_body, draw_hand=include_hand, draw_face=include_face) detected_map = HWC3(detected_map) img = resize_image(input_image, image_resolution) From 064200ecf52a6fd913ce00c6ee1fd81c20b92b3c Mon Sep 17 00:00:00 2001 From: mymusise Date: Wed, 17 Jul 2024 17:45:39 +0800 Subject: [PATCH 2/2] add hand_and_face && change default value Signed-off-by: mymusise --- src/controlnet_aux/dwpose/__init__.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/controlnet_aux/dwpose/__init__.py b/src/controlnet_aux/dwpose/__init__.py index f009223..ad6aafc 100644 --- a/src/controlnet_aux/dwpose/__init__.py +++ b/src/controlnet_aux/dwpose/__init__.py @@ -43,8 +43,11 @@ def to(self, device): self.pose_estimation.to(device) return self - def __call__(self, input_image, detect_resolution=512, image_resolution=512, include_body=True, include_hand=False, include_face=False, output_type="pil", **kwargs): - + def __call__(self, input_image, detect_resolution=512, image_resolution=512, include_body=True, include_hand=True, include_face=True, hand_and_face=None, output_type="pil", **kwargs): + if hand_and_face is not None: + include_hand = hand_and_face + include_face = hand_and_face + input_image = cv2.cvtColor(np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR) input_image = HWC3(input_image)