From cdac68a9a9b50162cb59849c8670ca72aba2b8ad Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Tue, 30 Jan 2024 13:30:43 +0800
Subject: [PATCH] add reverse mask

---
 iopaint/helper.py                             | 32 +++++++++++--------
 iopaint/schema.py                             |  4 +--
 iopaint/tests/test_adjust_mask.py             |  2 ++
 .../components/SidePanel/DiffusionOptions.tsx | 11 +++++++
 web_app/src/lib/types.ts                      |  2 +-
 5 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/iopaint/helper.py b/iopaint/helper.py
index ce83f38..d0a9943 100644
--- a/iopaint/helper.py
+++ b/iopaint/helper.py
@@ -354,25 +354,29 @@ def concat_alpha_channel(rgb_np_img, alpha_channel) -> np.ndarray:
 
 
 def adjust_mask(mask: np.ndarray, kernel_size: int, operate):
+    # fronted brush color "ffcc00bb"
     # kernel_size = kernel_size*2+1
     mask[mask >= 127] = 255
     mask[mask < 127] = 0
-    # fronted brush color "ffcc00bb"
-    kernel = cv2.getStructuringElement(
-        cv2.MORPH_ELLIPSE, (2 * kernel_size + 1, 2 * kernel_size + 1)
-    )
-    if operate == "expand":
-        mask = cv2.dilate(
-            mask,
-            kernel,
-            iterations=1,
-        )
+
+    if operate == "reverse":
+        mask = 255 - mask
     else:
-        mask = cv2.erode(
-            mask,
-            kernel,
-            iterations=1,
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, (2 * kernel_size + 1, 2 * kernel_size + 1)
         )
+        if operate == "expand":
+            mask = cv2.dilate(
+                mask,
+                kernel,
+                iterations=1,
+            )
+        else:
+            mask = cv2.erode(
+                mask,
+                kernel,
+                iterations=1,
+            )
     res_mask = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
     res_mask[mask > 128] = [255, 203, 0, int(255 * 0.73)]
     res_mask = cv2.cvtColor(res_mask, cv2.COLOR_BGRA2RGBA)
diff --git a/iopaint/schema.py b/iopaint/schema.py
index 739a07f..f1f82d7 100644
--- a/iopaint/schema.py
+++ b/iopaint/schema.py
@@ -304,12 +304,12 @@ class SwitchModelRequest(BaseModel):
     name: str
 
 
-AdjustMaskOperate = Literal["expand", "shrink"]
+AdjustMaskOperate = Literal["expand", "shrink", "reverse"]
 
 
 class AdjustMaskRequest(BaseModel):
     mask: str = Field(
         ..., description="base64 encoded mask. 255 means area to do inpaint"
     )
-    operate: AdjustMaskOperate = Field(..., description="expand or shrink")
+    operate: AdjustMaskOperate = Field(..., description="expand/shrink/reverse")
     kernel_size: int = Field(5, description="Kernel size for expanding mask")
diff --git a/iopaint/tests/test_adjust_mask.py b/iopaint/tests/test_adjust_mask.py
index dfa8040..1f01713 100644
--- a/iopaint/tests/test_adjust_mask.py
+++ b/iopaint/tests/test_adjust_mask.py
@@ -13,3 +13,5 @@ def test_adjust_mask():
     cv2.imwrite(str(save_dir / "adjust_mask_expand.png"), res_mask)
     res_mask = adjust_mask(mask, 20, "shrink")
     cv2.imwrite(str(save_dir / "adjust_mask_shrink.png"), res_mask)
+    res_mask = adjust_mask(mask, 20, "reverse")
+    cv2.imwrite(str(save_dir / "adjust_mask_reverse.png"), res_mask)
diff --git a/web_app/src/components/SidePanel/DiffusionOptions.tsx b/web_app/src/components/SidePanel/DiffusionOptions.tsx
index 19d6820..6687274 100644
--- a/web_app/src/components/SidePanel/DiffusionOptions.tsx
+++ b/web_app/src/components/SidePanel/DiffusionOptions.tsx
@@ -839,6 +839,17 @@ const DiffusionOptions = () => {
                   Shrink
                 </div>
               </Button>
+
+              <Button
+                variant="outline"
+                className="p-1 h-8"
+                onClick={() => adjustMask("reverse")}
+                disabled={isProcessing}
+              >
+                <div className="flex items-center gap-1 select-none">
+                  Reverse
+                </div>
+              </Button>
             </div>
 
             <Button
diff --git a/web_app/src/lib/types.ts b/web_app/src/lib/types.ts
index 9d96008..0c2383e 100644
--- a/web_app/src/lib/types.ts
+++ b/web_app/src/lib/types.ts
@@ -126,4 +126,4 @@ export enum PowerPaintTask {
   outpainting = "outpainting",
 }
 
-export type AdjustMaskOperate = "expand" | "shrink"
+export type AdjustMaskOperate = "expand" | "shrink" | "reverse"