Merge pull request huggingface#49 from pdoane/annotator-params

Consistent resizing parameters and behavior

Author: patrickvonplaten

.gitignore

@@ -13,6 +13,7 @@ tests/fixtures/cached_*_text.txt
 logs/
 lightning_logs/
 lang_code_data/
+tests/outputs
 
 # Distribution / packaging
 .Python

src/controlnet_aux/canny/__init__.py

@@ -1,21 +1,36 @@
+import warnings
 import cv2
 import numpy as np
 from PIL import Image
-from ..util import HWC3
+from ..util import HWC3, resize_image
 
 class CannyDetector:
-    def __call__(self, img, low_threshold=100, high_threshold=200):
+    def __call__(self, input_image=None, low_threshold=100, high_threshold=200, detect_resolution=512, image_resolution=512, output_type=None, **kwargs):
+        if "img" in kwargs:
+            warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+            input_image = kwargs.pop("img")
 
-        input_type = "np"
-        if isinstance(img, Image.Image):
-            img = np.array(img)
-            input_type = "pil"
+        if input_image is None:
+            raise ValueError("input_image must be defined.")
+
+        if not isinstance(input_image, np.ndarray):
+            input_image = np.array(input_image, dtype=np.uint8)
+            output_type = output_type or "pil"
+        else:
+            output_type = output_type or "np"
 
-        img = HWC3(img)
-        img = cv2.Canny(img, low_threshold, high_threshold)
+        input_image = HWC3(input_image)
+        input_image = resize_image(input_image, detect_resolution)
+
+        detected_map = cv2.Canny(input_image, low_threshold, high_threshold)
+        detected_map = HWC3(detected_map)      
+         
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
 
-        if input_type == "pil":
-            img = Image.fromarray(img)
-            img = img.convert("RGB")
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
 
-        return img
+        return detected_map

src/controlnet_aux/hed/__init__.py

@@ -6,6 +6,7 @@
 # and in this way it works better for gradio's RGB protocol
 
 import os
+import warnings
 
 import cv2
 import numpy as np
@@ -78,7 +79,15 @@ def to(self, device):
         self.netNetwork.to(device)
         return self
 
-    def __call__(self, input_image, detect_resolution=512, image_resolution=512, safe=False, return_pil=True, scribble=False):
+    def __call__(self, input_image, detect_resolution=512, image_resolution=512, safe=False, output_type="pil", scribble=False, **kwargs):
+        if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+        if type(output_type) is bool:
+            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+            if output_type:
+                output_type = "pil"
+
         device = next(iter(self.netNetwork.parameters())).device
         if not isinstance(input_image, np.ndarray):
             input_image = np.array(input_image, dtype=np.uint8)
@@ -101,7 +110,6 @@ def __call__(self, input_image, detect_resolution=512, image_resolution=512, saf
             edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
 
         detected_map = edge
-
         detected_map = HWC3(detected_map)
 
         img = resize_image(input_image, image_resolution)
@@ -115,7 +123,7 @@ def __call__(self, input_image, detect_resolution=512, image_resolution=512, saf
             detected_map[detected_map > 4] = 255
             detected_map[detected_map < 255] = 0
 
-        if return_pil:
+        if output_type == "pil":
             detected_map = Image.fromarray(detected_map)
 
         return detected_map

src/controlnet_aux/lineart/__init__.py

@@ -1,4 +1,5 @@
 import os
+import warnings
 
 import cv2
 import numpy as np
@@ -122,7 +123,15 @@ def to(self, device):
         self.model_coarse.to(device)
         return self
 
-    def __call__(self, input_image, coarse=False, detect_resolution=512, image_resolution=512, return_pil=True):
+    def __call__(self, input_image, coarse=False, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
+        if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+        if type(output_type) is bool:
+            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+            if output_type:
+                output_type = "pil"
+
         device = next(iter(self.model.parameters())).device
         if not isinstance(input_image, np.ndarray):
             input_image = np.array(input_image, dtype=np.uint8)
@@ -152,7 +161,7 @@ def __call__(self, input_image, coarse=False, detect_resolution=512, image_resol
         detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
         detected_map = 255 - detected_map
 
-        if return_pil:
+        if output_type == "pil":
             detected_map = Image.fromarray(detected_map)
 
         return detected_map

src/controlnet_aux/lineart_anime/__init__.py

@@ -1,10 +1,11 @@
-import numpy as np
-import torch
-import torch.nn as nn
 import functools
-
 import os
+import warnings
+
 import cv2
+import numpy as np
+import torch
+import torch.nn as nn
 from einops import rearrange
 from huggingface_hub import hf_hub_download
 from PIL import Image
@@ -141,7 +142,15 @@ def to(self, device):
         self.model.to(device)
         return self
 
-    def __call__(self, input_image, detect_resolution=512, image_resolution=512, return_pil=True):
+    def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
+        if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+        if type(output_type) is bool:
+            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+            if output_type:
+                output_type = "pil"
+
         device = next(iter(self.model.parameters())).device
         if not isinstance(input_image, np.ndarray):
             input_image = np.array(input_image, dtype=np.uint8)
@@ -174,7 +183,7 @@ def __call__(self, input_image, detect_resolution=512, image_resolution=512, ret
         detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
         detected_map = 255 - detected_map
 
-        if return_pil:
+        if output_type == "pil":
             detected_map = Image.fromarray(detected_map)
 
         return detected_map

src/controlnet_aux/mediapipe_face/__init__.py

@@ -1,23 +1,53 @@
+import warnings
 from typing import Union
 
-from .mediapipe_face_common import generate_annotation
-from PIL import Image
+import cv2
 import numpy as np
+from PIL import Image
+
+from ..util import HWC3, resize_image
+from .mediapipe_face_common import generate_annotation
 
 
 class MediapipeFaceDetector:
     def __call__(self,
-                 image: Union[np.ndarray, Image.Image],
+                 input_image: Union[np.ndarray, Image.Image] = None,
                  max_faces: int = 1,
                  min_confidence: float = 0.5,
-                 return_pil: bool = True):
-        
-        if isinstance(image, Image.Image) is True:
-            image = np.array(image)
-        
-        face = generate_annotation(image, max_faces, min_confidence)
+                 output_type: str = "pil",
+                 detect_resolution: int = 512,
+                 image_resolution: int = 512,
+                 **kwargs):
+
+        if "image" in kwargs:
+            warnings.warn("image is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+            input_image = kwargs.pop("image")        
+        if input_image is None:
+            raise ValueError("input_image must be defined.")
+
+        if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+        if type(output_type) is bool:
+            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+            if output_type:
+                output_type = "pil"
+
+        if not isinstance(input_image, np.ndarray):
+            input_image = np.array(input_image, dtype=np.uint8)
+
+        input_image = HWC3(input_image)
+        input_image = resize_image(input_image, detect_resolution)
+
+        detected_map = generate_annotation(input_image, max_faces, min_confidence)
+        detected_map = HWC3(detected_map)
+
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
 
-        if return_pil is True:
-            face = Image.fromarray(face)
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
 
-        return face
+        return detected_map

src/controlnet_aux/midas/__init__.py

@@ -7,7 +7,7 @@
 from huggingface_hub import hf_hub_download
 from PIL import Image
 
-from ..util import HWC3
+from ..util import HWC3, resize_image
 from .api import MiDaSInference
 
 
@@ -36,14 +36,17 @@ def to(self, device):
         self.model.to(device)
         return self
 
-    def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1, depth_and_normal=False):
+    def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1, depth_and_normal=False, detect_resolution=512, image_resolution=512, output_type=None):
         device = next(iter(self.model.parameters())).device
-        input_type = "np"
-        if isinstance(input_image, Image.Image):
-            input_image = np.array(input_image)
-            input_type = "pil"
-
+        if not isinstance(input_image, np.ndarray):
+            input_image = np.array(input_image, dtype=np.uint8)
+            output_type = output_type or "pil"
+        else:
+            output_type = output_type or "np"
+        
         input_image = HWC3(input_image)
+        input_image = resize_image(input_image, detect_resolution)
+
         assert input_image.ndim == 3
         image_depth = input_image
         with torch.no_grad():
@@ -70,9 +73,19 @@ def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1, depth_and_normal=False
                 normal /= np.sum(normal ** 2.0, axis=2, keepdims=True) ** 0.5
                 normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)[:, :, ::-1]
 
-        if input_type == "pil":
+        depth_image = HWC3(depth_image)
+        if depth_and_normal:
+            normal_image = HWC3(normal_image)
+
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        depth_image = cv2.resize(depth_image, (W, H), interpolation=cv2.INTER_LINEAR)
+        if depth_and_normal:
+            normal_image = cv2.resize(normal_image, (W, H), interpolation=cv2.INTER_LINEAR)
+        
+        if output_type == "pil":
             depth_image = Image.fromarray(depth_image)
-            depth_image = depth_image.convert("RGB")
             if depth_and_normal:
                 normal_image = Image.fromarray(normal_image)
 

src/controlnet_aux/mlsd/__init__.py

@@ -1,9 +1,9 @@
 import os
+import warnings
 
 import cv2
 import numpy as np
 import torch
-from einops import rearrange
 from huggingface_hub import hf_hub_download
 from PIL import Image
 
@@ -38,7 +38,15 @@ def to(self, device):
         self.model.to(device)
         return self
 
-    def __call__(self, input_image, thr_v=0.1, thr_d=0.1, detect_resolution=512, image_resolution=512, return_pil=True):
+    def __call__(self, input_image, thr_v=0.1, thr_d=0.1, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
+        if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+        if type(output_type) is bool:
+            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+            if output_type:
+                output_type = "pil"
+
         if not isinstance(input_image, np.ndarray):
             input_image = np.array(input_image, dtype=np.uint8)
 
@@ -58,14 +66,14 @@ def __call__(self, input_image, thr_v=0.1, thr_d=0.1, detect_resolution=512, ima
             pass
 
         detected_map = img_output[:, :, 0]
-
         detected_map = HWC3(detected_map)
+
         img = resize_image(input_image, image_resolution)
         H, W, C = img.shape
 
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST)
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
 
-        if return_pil:
+        if output_type == "pil":
             detected_map = Image.fromarray(detected_map)
 
         return detected_map

src/controlnet_aux/normalbae/__init__.py

@@ -1,6 +1,8 @@
 import os
 import types
+import warnings
 
+import cv2
 import numpy as np
 import torch
 import torchvision.transforms as transforms
@@ -58,7 +60,15 @@ def to(self, device):
         return self
 
 
-    def __call__(self, input_image, detect_resolution=512, image_resolution=512, return_pil=True):
+    def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
+        if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+        if type(output_type) is bool:
+            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+            if output_type:
+                output_type = "pil"
+
         device = next(iter(self.model.parameters())).device
         if not isinstance(input_image, np.ndarray):
             input_image = np.array(input_image, dtype=np.uint8)
@@ -84,10 +94,16 @@ def __call__(self, input_image, detect_resolution=512, image_resolution=512, ret
             normal = rearrange(normal[0], 'c h w -> h w c').cpu().numpy()
             normal_image = (normal * 255.0).clip(0, 255).astype(np.uint8)
 
-        img = resize_image(normal_image, image_resolution)
+        detected_map = normal_image
+        detected_map = HWC3(detected_map)      
+
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
 
-        if return_pil:
-            img = Image.fromarray(img)
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
 
-        return img
+        return detected_map