From 019cd099272898764d709d1799292664d52d033a Mon Sep 17 00:00:00 2001 From: Bingsu Date: Sun, 7 May 2023 13:39:05 +0900 Subject: [PATCH] refactor: mask preprocess, save image --- adetailer/common.py | 37 ++++++++++++++++++++++++++ scripts/!adetailer.py | 60 ++++++++++++++++++++++++++----------------- 2 files changed, 73 insertions(+), 24 deletions(-) diff --git a/adetailer/common.py b/adetailer/common.py index 46bd20d..731b081 100644 --- a/adetailer/common.py +++ b/adetailer/common.py @@ -137,3 +137,40 @@ def offset(img: Image.Image, x: int = 0, y: int = 0) -> Image.Image: def is_all_black(img: Image.Image) -> bool: arr = np.array(img) return cv2.countNonZero(arr) == 0 + + +def mask_preprocess( + masks: list[Image.Image] | None, + kernel: int = 0, + x_offset: int = 0, + y_offset: int = 0, +) -> list[Image.Image]: + """ + The mask_preprocess function takes a list of masks and preprocesses them. + It dilates and erodes the masks, and offsets them by x_offset and y_offset. + + Parameters + ---------- + masks: list[Image.Image] | None + A list of masks + kernel: int + kernel size of dilation or erosion + x_offset: int + → + y_offset: int + ↑ + + Returns + ------- + list[Image.Image] + A list of processed masks + """ + if masks is None: + return [] + + masks = [dilate_erode(m, kernel) for m in masks] + masks = [m for m in masks if not is_all_black(m)] + if x_offset != 0 or y_offset != 0: + masks = [offset(m, x_offset, y_offset) for m in masks] + + return masks diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py index d3c5d26..1c4bb02 100644 --- a/scripts/!adetailer.py +++ b/scripts/!adetailer.py @@ -22,7 +22,7 @@ from adetailer import ( mediapipe_predict, ultralytics_predict, ) -from adetailer.common import dilate_erode, is_all_black, offset +from adetailer.common import mask_preprocess from controlnet_ext import ControlNetExt, controlnet_exists, get_cn_inpaint_models from modules import images, safe, script_callbacks, scripts, shared from modules.paths import data_path, models_path @@ -277,6 +277,9 @@ class AfterDetailerScript(scripts.Script): ) def update_controlnet_args(self, p, args: ADetailerArgs) -> None: + if self.controlnet_ext is None: + self.init_controlnet_ext() + if ( self.controlnet_ext is not None and self.controlnet_ext.cn_available @@ -288,8 +291,9 @@ class AfterDetailerScript(scripts.Script): def is_ad_enabled(self, *args_) -> bool: if len(args_) < 2: - message = f"""[-] ADetailer: Not enough arguments passed to adetailer. - input: {args_!r} + message = f""" + [-] ADetailer: Not enough arguments passed to adetailer. + input: {args_!r} """ raise ValueError(dedent(message)) checker = EnableChecker(ad_enable=args_[0], ad_model=args_[1]) @@ -447,12 +451,13 @@ class AfterDetailerScript(scripts.Script): i2i._disable_adetailer = True if args.ad_controlnet_model != "None": - self.init_controlnet_ext() self.update_controlnet_args(i2i, args) return i2i - def save_image(self, p, image, seed, *, condition: str, suffix: str) -> None: + def save_image(self, p, image, *, condition: str, suffix: str) -> None: i = p._idx + seed, _ = self.get_seed(p) + if opts.data.get(condition, False): images.save_image( image=image, @@ -481,17 +486,20 @@ class AfterDetailerScript(scripts.Script): extra_params = self.extra_params(args) p.extra_generation_params.update(extra_params) - def _postprocess_image(self, p, pp, args: ADetailerArgs): + def _postprocess_image(self, p, pp, args: ADetailerArgs) -> bool: + """ + Returns + ------- + bool + + `True` if image was processed, `False` otherwise. + """ p._idx = getattr(p, "_idx", -1) + 1 i = p._idx i2i = self.get_i2i_p(p, args, pp.image) seed, subseed = self.get_seed(p) - self.save_image( - p, pp.image, seed, condition="ad_save_images_before", suffix="-ad-before" - ) - is_mediapipe = args.ad_model.lower().startswith("mediapipe") kwargs = {} @@ -506,17 +514,18 @@ class AfterDetailerScript(scripts.Script): with ChangeTorchLoad(): pred = predictor(ad_model, pp.image, args.ad_conf, **kwargs) - if pred.masks is None: + masks = mask_preprocess(pred.masks) + + if not masks: print( f"[-] ADetailer: nothing detected on image {i + 1} with current settings." ) - return + return False self.save_image( - p, pred.preview, seed, condition="ad_save_previews", suffix="-ad-preview" + p, pred.preview, condition="ad_save_previews", suffix="-ad-preview" ) - masks = pred.masks steps = len(masks) processed = None @@ -525,22 +534,20 @@ class AfterDetailerScript(scripts.Script): p2 = copy(i2i) for j in range(steps): - mask = masks[j] - mask = dilate_erode(mask, args.ad_dilate_erode) + p2.image_mask = masks[j] + processed = process_images(p2) - if not is_all_black(mask): - mask = offset(mask, args.ad_x_offset, args.ad_y_offset) - p2.image_mask = mask - processed = process_images(p2) - - p2 = copy(i2i) - p2.init_images = [processed.images[0]] + p2 = copy(i2i) + p2.init_images = [processed.images[0]] p2.seed = seed + j + 1 p2.subseed = subseed + j + 1 if processed is not None: pp.image = processed.images[0] + return True + + return False def postprocess_image(self, p, pp, *args_): if getattr(p, "_disable_adetailer", False): @@ -550,7 +557,12 @@ class AfterDetailerScript(scripts.Script): return args = self.get_args(*args_) - self._postprocess_image(p, pp, args) + is_processed = self._postprocess_image(p, pp, args) + + if is_processed: + self.save_image( + p, pp.image, condition="ad_save_images_before", suffix="-ad-before" + ) try: if p._idx == len(p.all_prompts) - 1: