From 2a362ec845ef2df81634b29c3c2a70e4fa9f3d6e Mon Sep 17 00:00:00 2001
From: Bingsu <ks2515@naver.com>
Date: Tue, 8 Aug 2023 20:21:11 +0900
Subject: [PATCH] fix: pr style

---
 adetailer/args.py     | 5 +++--
 adetailer/mask.py     | 2 +-
 adetailer/ui.py       | 2 +-
 scripts/!adetailer.py | 4 ++--
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/adetailer/args.py b/adetailer/args.py
index 03ebdf4..1a0ac93 100644
--- a/adetailer/args.py
+++ b/adetailer/args.py
@@ -41,6 +41,7 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
     ad_prompt: str = ""
     ad_negative_prompt: str = ""
     ad_confidence: confloat(ge=0.0, le=1.0) = 0.3
+    ad_mask_k_largest: NonNegativeInt = 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 = 4
@@ -70,7 +71,6 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
     ad_controlnet_weight: confloat(ge=0.0, le=1.0) = 1.0
     ad_controlnet_guidance_start: confloat(ge=0.0, le=1.0) = 0.0
     ad_controlnet_guidance_end: confloat(ge=0.0, le=1.0) = 1.0
-    ad_mask_k_largest: NonNegativeInt = 0
     is_api: bool = True
 
     @root_validator(skip_on_failure=True)
@@ -113,6 +113,7 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
 
         ppop("ADetailer prompt")
         ppop("ADetailer negative prompt")
+        ppop("ADetailer mask only top k largest", cond=0)
         ppop("ADetailer mask min ratio", cond=0.0)
         ppop("ADetailer mask max ratio", cond=1.0)
         ppop("ADetailer x offset", cond=0)
@@ -189,9 +190,9 @@ _all_args = [
     ("ad_prompt", "ADetailer prompt"),
     ("ad_negative_prompt", "ADetailer negative prompt"),
     ("ad_confidence", "ADetailer confidence"),
+    ("ad_mask_k_largest", "ADetailer mask only top k largest"),
     ("ad_mask_min_ratio", "ADetailer mask min ratio"),
     ("ad_mask_max_ratio", "ADetailer mask max ratio"),
-    ("ad_mask_k_largest", "ADetailer mask only top k largest"),
     ("ad_x_offset", "ADetailer x offset"),
     ("ad_y_offset", "ADetailer y offset"),
     ("ad_dilate_erode", "ADetailer dilate/erode"),
diff --git a/adetailer/mask.py b/adetailer/mask.py
index eaa90c8..913ea7e 100644
--- a/adetailer/mask.py
+++ b/adetailer/mask.py
@@ -215,7 +215,7 @@ def filter_by_ratio(pred: PredictOutput, low: float, high: float) -> PredictOutp
     return pred
 
 
-def filter_take_largest(pred: PredictOutput, k: int) -> PredictOutput:
+def filter_k_largest(pred: PredictOutput, k: int = 0) -> PredictOutput:
     if not pred.bboxes or k == 0:
         return pred
     areas = [bbox_area(bbox) for bbox in pred.bboxes]
diff --git a/adetailer/ui.py b/adetailer/ui.py
index 6b78f9e..7b6261e 100644
--- a/adetailer/ui.py
+++ b/adetailer/ui.py
@@ -207,7 +207,7 @@ def detection(w: Widgets, n: int, is_img2img: bool):
             w.ad_mask_k_largest = gr.Slider(
                 label="Mask only the top k largest (0 to disable)" + suffix(n),
                 minumum=0,
-                maximum=5,
+                maximum=10,
                 step=1,
                 value=0,
                 visible=True,
diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py
index e599ceb..cf7016a 100644
--- a/scripts/!adetailer.py
+++ b/scripts/!adetailer.py
@@ -28,7 +28,7 @@ from adetailer.args import ALL_ARGS, BBOX_SORTBY, ADetailerArgs, EnableChecker
 from adetailer.common import PredictOutput
 from adetailer.mask import (
     filter_by_ratio,
-    filter_take_largest,
+    filter_k_largest,
     mask_preprocess,
     sort_bboxes,
 )
@@ -468,7 +468,7 @@ class AfterDetailerScript(scripts.Script):
         pred = filter_by_ratio(
             pred, low=args.ad_mask_min_ratio, high=args.ad_mask_max_ratio
         )
-        pred = filter_take_largest(pred, k=args.ad_mask_k_largest)
+        pred = filter_k_largest(pred, k=args.ad_mask_k_largest)
         pred = self.sort_bboxes(pred)
         return mask_preprocess(
             pred.masks,