update

Author: zhouzhengguang

src/controlnet_aux/dwpose/__init__.py

@@ -12,6 +12,7 @@
 import numpy as np
 from PIL import Image
 
+from ..util import HWC3, resize_image
 from . import util
 from .wholebody import Wholebody
 
@@ -31,17 +32,16 @@ def draw_pose(pose, H, W):
     return canvas
 
 class DWposeDetector:
-    def __init__(self, det_config, det_ckpt, pose_config, pose_ckpt, device):
+    def __init__(self, det_config=None, det_ckpt=None, pose_config=None, pose_ckpt=None, device="cpu"):
 
         self.pose_estimation = Wholebody(det_config, det_ckpt, pose_config, pose_ckpt, device)
 
-    def __call__(self, oriImg, output_type="pil", detect_resolution=512, image_resolution=512):
+    def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
 
-        oriImg = oriImg.copy()
-        input_image = cv2.cvtColor(np.array(oriImg), cv2.COLOR_RGB2BGR)
+        input_image = cv2.cvtColor(np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR)
 
-        input_image = util.HWC3(input_image)
-        input_image = util.resize_image(input_image, detect_resolution)
+        input_image = HWC3(input_image)
+        input_image = resize_image(input_image, detect_resolution)
         H, W, C = input_image.shape
 
         with torch.no_grad():
@@ -74,9 +74,9 @@ def __call__(self, oriImg, output_type="pil", detect_resolution=512, image_resol
             pose = dict(bodies=bodies, hands=hands, faces=faces)
 
             detected_map = draw_pose(pose, H, W)
-            detected_map = util.HWC3(detected_map)
+            detected_map = HWC3(detected_map)
 
-            img = util.resize_image(input_image, image_resolution)
+            img = resize_image(input_image, image_resolution)
             H, W, C = img.shape
 
             detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)

src/controlnet_aux/dwpose/util.py

@@ -7,37 +7,6 @@
 eps = 0.01
 
 
-def HWC3(x):
-    assert x.dtype == np.uint8
-    if x.ndim == 2:
-        x = x[:, :, None]
-    assert x.ndim == 3
-    H, W, C = x.shape
-    assert C == 1 or C == 3 or C == 4
-    if C == 3:
-        return x
-    if C == 1:
-        return np.concatenate([x, x, x], axis=2)
-    if C == 4:
-        color = x[:, :, 0:3].astype(np.float32)
-        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
-        y = color * alpha + 255.0 * (1.0 - alpha)
-        y = y.clip(0, 255).astype(np.uint8)
-        return y
-
-def resize_image(input_image, resolution):
-    H, W, C = input_image.shape
-    H = float(H)
-    W = float(W)
-    k = float(resolution) / min(H, W)
-    H *= k
-    W *= k
-    H = int(np.round(H / 64.0)) * 64
-    W = int(np.round(W / 64.0)) * 64
-    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
-    return img
-
-
 def smart_resize(x, s):
     Ht, Wt = s
     if x.ndim == 2:

src/controlnet_aux/dwpose/wholebody.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import os
 import numpy as np
 from . import util
 import cv2
@@ -19,7 +20,13 @@ def __init__(self,
                  det_config=None, det_ckpt=None, 
                  pose_config=None, pose_ckpt=None,
                  device="cpu"):
-                
+        
+        if det_config is None:
+            det_config = os.path.join(os.path.dirname(__file__), "dwpose_config/dwpose-l_384x288.py")
+        
+        if pose_config is None:
+            pose_config = os.path.join(os.path.dirname(__file__), "yolox_config/yolox_l_8xb8-300e_coco.py")
+
         if det_ckpt is None:
             det_ckpt = 'https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_l_8x8_300e_coco/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth'
 

tests/test_controlnet_aux.py

@@ -120,8 +120,7 @@ def test_zoe(img):
     zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
     common("zoe", zoe, img)
 
-def test_dwpose(img, det_config, det_ckpt, pose_config, pose_ckpt):
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    dwpose = DWposeDetector(det_config, det_ckpt, pose_config, pose_ckpt, device)
+def test_dwpose(img):
+    dwpose = DWposeDetector()
     common("dwpose", dwpose, img)
     return_pil("dwpose", dwpose, img)