From d9d89f32b76ea539015c41c8df028b58a5e9b9bd Mon Sep 17 00:00:00 2001
From: Bingsu <ks2515@naver.com>
Date: Thu, 18 May 2023 21:53:39 +0900
Subject: [PATCH] fix: filter by ratio step, mask merge invert

---
 adetailer/__version__.py |  2 +-
 adetailer/args.py        |  6 +++---
 adetailer/mask.py        | 38 +++++++++++++++++++++++++++++++++++++-
 scripts/!adetailer.py    | 23 +++++++++++++++--------
 4 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/adetailer/__version__.py b/adetailer/__version__.py
index 35e4003..72f0fe5 100644
--- a/adetailer/__version__.py
+++ b/adetailer/__version__.py
@@ -1 +1 @@
-__version__ = "23.5.15"
+__version__ = "23.5.16.dev0"
diff --git a/adetailer/args.py b/adetailer/args.py
index b331a2f..e146228 100644
--- a/adetailer/args.py
+++ b/adetailer/args.py
@@ -37,8 +37,8 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
     ad_prompt: str = ""
     ad_negative_prompt: str = ""
     ad_conf: confloat(ge=0.0, le=1.0) = 0.3
-    ad_mask_min_scale: confloat(ge=0.0, le=1.0) = 0.0
-    ad_mask_max_scale: confloat(ge=0.0, le=1.0) = 1.0
+    ad_mask_min_ratio: confloat(ge=0.0, le=1.0) = 0.0
+    ad_mask_max_ratio: confloat(ge=0.0, le=1.0) = 1.0
     ad_dilate_erode: int = 32
     ad_x_offset: int = 0
     ad_y_offset: int = 0
@@ -100,7 +100,7 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
         ppop("ADetailer mask max ratio", cond=1.0)
         ppop("ADetailer x offset", cond=0)
         ppop("ADetailer y offset", cond=0)
-        ppop("ad_mask_merge_invert", cond=0)
+        ppop("ADetailer mask merge/invert", cond=0)
         ppop("ADetailer inpaint full", ["ADetailer inpaint padding"])
         ppop(
             "ADetailer use inpaint width/height",
diff --git a/adetailer/mask.py b/adetailer/mask.py
index 3860a65..d386b85 100644
--- a/adetailer/mask.py
+++ b/adetailer/mask.py
@@ -1,5 +1,5 @@
 from enum import IntEnum
-from functools import partial
+from functools import partial, reduce
 from math import dist
 
 import cv2
@@ -16,6 +16,12 @@ class SortBy(IntEnum):
     AREA = 3
 
 
+class MergeInvert(IntEnum):
+    NONE = 0
+    MERGE = 1
+    MERGE_INVERT = 2
+
+
 def _dilate(arr: np.ndarray, value: int) -> np.ndarray:
     kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
     return cv2.dilate(arr, kernel, iterations=1)
@@ -88,6 +94,7 @@ def mask_preprocess(
     kernel: int = 0,
     x_offset: int = 0,
     y_offset: int = 0,
+    merge_invert: int | MergeInvert = MergeInvert.NONE,
 ) -> list[Image.Image]:
     """
     The mask_preprocess function takes a list of masks and preprocesses them.
@@ -119,6 +126,8 @@ def mask_preprocess(
         masks = [dilate_erode(m, kernel) for m in masks]
         masks = [m for m in masks if not is_all_black(m)]
 
+    masks = mask_merge_invert(masks, mode=merge_invert)
+
     return masks
 
 
@@ -203,3 +212,30 @@ def filter_by_ratio(pred: PredictOutput, low: float, high: float) -> PredictOutp
     pred.bboxes = [pred.bboxes[i] for i in idx]
     pred.masks = [pred.masks[i] for i in idx]
     return pred
+
+
+# Merge / Invert
+def mask_merge(masks: list[Image.Image]) -> list[Image.Image]:
+    arrs = [np.array(m) for m in masks]
+    arr = reduce(cv2.bitwise_or, arrs)
+    return [Image.fromarray(arr)]
+
+
+def mask_invert(masks: list[Image.Image]) -> list[Image.Image]:
+    return [ImageChops.invert(m) for m in masks]
+
+
+def mask_merge_invert(
+    masks: list[Image.Image], mode: int | MergeInvert
+) -> list[Image.Image]:
+    if mode == MergeInvert.NONE or not masks:
+        return masks
+
+    if mode == MergeInvert.MERGE:
+        return mask_merge(masks)
+
+    if mode == MergeInvert.MERGE_INVERT:
+        merged = mask_merge(masks)
+        return mask_invert(merged)
+
+    raise RuntimeError
diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py
index 4ce708d..f9495fa 100644
--- a/scripts/!adetailer.py
+++ b/scripts/!adetailer.py
@@ -23,7 +23,7 @@ from adetailer import (
 )
 from adetailer.args import ALL_ARGS, BBOX_SORTBY, ADetailerArgs, EnableChecker
 from adetailer.common import PredictOutput
-from adetailer.mask import mask_preprocess, sort_bboxes
+from adetailer.mask import filter_by_ratio, mask_preprocess, sort_bboxes
 from adetailer.ui import adui, ordinal, suffix
 from controlnet_ext import ControlNetExt, controlnet_exists
 from sd_webui import images, safe, script_callbacks, scripts, shared
@@ -384,6 +384,19 @@ class AfterDetailerScript(scripts.Script):
         pred = sort_bboxes(pred, sortby_idx)
         return pred
 
+    def pred_preprocessing(self, pred: PredictOutput, args: ADetailerArgs):
+        pred = filter_by_ratio(
+            pred, low=args.ad_mask_min_ratio, high=args.ad_mask_max_ratio
+        )
+        pred = self.sort_bboxes(pred)
+        return mask_preprocess(
+            pred.masks,
+            kernel=args.ad_dilate_erode,
+            x_offset=args.ad_x_offset,
+            y_offset=args.ad_y_offset,
+            merge_invert=args.ad_mask_merge_invert,
+        )
+
     def i2i_prompts_replace(
         self, i2i, prompts: list[str], negative_prompts: list[str], j: int
     ):
@@ -431,13 +444,7 @@ class AfterDetailerScript(scripts.Script):
         with ChangeTorchLoad():
             pred = predictor(ad_model, pp.image, args.ad_conf, **kwargs)
 
-        pred = self.sort_bboxes(pred)
-        masks = mask_preprocess(
-            pred.masks,
-            kernel=args.ad_dilate_erode,
-            x_offset=args.ad_x_offset,
-            y_offset=args.ad_y_offset,
-        )
+        masks = self.pred_preprocessing(pred, args)
 
         if not masks:
             print(