Merge pull request Mikubill#495 from KurozumiGH/preproc_binary

Add binary preprocessor

annotator/binary/__init__.py

@@ -0,0 +1,14 @@
+import cv2
+
+
+def apply_binary(img, bin_threshold):
+    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+
+    if bin_threshold == 0 or bin_threshold == 255:
+        # Otsu's threshold
+        otsu_threshold, img_bin = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+        print("Otsu threshold:", otsu_threshold)
+    else:
+        _, img_bin = cv2.threshold(img_gray, bin_threshold, 255, cv2.THRESH_BINARY_INV)
+
+    return cv2.cvtColor(img_bin, cv2.COLOR_GRAY2RGB)

scripts/api.py

@@ -272,7 +272,8 @@ async def detect(
             "mlsd", 
             "normal_map", 
             "openpose", 
-            "segmentation"
+            "segmentation",
+            "binary"
         ]
 
         if controlnet_module not in available_modules:
@@ -305,6 +306,8 @@ async def detect(
                 results.append(fake_scribble(img, controlnet_processor_res))
             elif controlnet_module == "segmentation":
                 results.append(uniformer(img, controlnet_processor_res))
+            elif controlnet_module == "binary":
+                results.append(binary(img, controlnet_processor_res, controlnet_threshold_a))
 
         if controlnet_module == "hed":
             unload_hed()

scripts/controlnet.py

@@ -192,6 +192,7 @@ def __init__(self) -> None:
             "scribble": simple_scribble,
             "fake_scribble": fake_scribble,
             "segmentation": uniformer,
+            "binary": binary,
         }
         self.unloadable = {
             "hed": unload_hed,
@@ -357,6 +358,13 @@ def build_sliders(module):
                     gr.update(label="Threshold B", value=64, minimum=64, maximum=1024, interactive=False),
                     gr.update(visible=True)
                 ]
+            elif module == "binary":
+                return [
+                    gr.update(label="Annotator resolution", value=512, minimum=64, maximum=2048, step=1, interactive=True),
+                    gr.update(label="Binary threshold", minimum=0, maximum=255, value=0, step=1, interactive=True),
+                    gr.update(label="Threshold B", value=64, minimum=64, maximum=1024, interactive=False),
+                    gr.update(visible=True)
+                ]
             elif module == "none":
                 return [
                     gr.update(label="Normal Resolution", value=64, minimum=64, maximum=2048, interactive=False),
@@ -793,7 +801,7 @@ def postprocess(self, p, processed, is_img2img=False, *args):
 
         if hasattr(self, "detected_map") and self.detected_map is not None:
             for detect_map, module in self.detected_map:
-                if module in ["canny", "mlsd", "scribble", "fake_scribble", "pidinet"]:
+                if module in ["canny", "mlsd", "scribble", "fake_scribble", "pidinet", "binary"]:
                     detect_map = 255-detect_map
                 processed.images.extend([Image.fromarray(detect_map)])
 

scripts/processor.py

@@ -212,4 +212,17 @@ def color(img, res=512, **kwargs):
         from annotator.color import apply_color
         model_color = apply_color
     result = model_color(img, res=res)
+    return result, True
+
+
+model_binary = None
+
+
+def binary(img, res=512, thr_a=0, **kwargs):
+    img = resize_image(HWC3(img), res)
+    global model_binary
+    if model_binary is None:
+        from annotator.binary import apply_binary
+        model_binary = apply_binary
+    result = model_binary(img, thr_a)
     return result, True