From acb8306916436cf4e82164c182f57f6f2b06877f Mon Sep 17 00:00:00 2001 From: Colin Redman <20376935+credman0@users.noreply.github.com> Date: Sun, 25 Dec 2022 22:40:29 -0600 Subject: [PATCH 1/4] Fix RGBA images throwing error in txt2mask --- shortcodes/stable_diffusion/txt2mask.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/shortcodes/stable_diffusion/txt2mask.py b/shortcodes/stable_diffusion/txt2mask.py index 81ade7d..9906e3e 100644 --- a/shortcodes/stable_diffusion/txt2mask.py +++ b/shortcodes/stable_diffusion/txt2mask.py @@ -14,6 +14,8 @@ def run_block(self, pargs, kwargs, context, content): from matplotlib import pyplot as plt import cv2 import numpy + from modules.images import flatten + from modules.shared import opts brush_mask_mode = self.Unprompted.parse_advanced(kwargs["mode"],context) if "mode" in kwargs else "add" self.show = True if "show" in pargs else False @@ -99,7 +101,8 @@ def get_mask(): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), transforms.Resize((512, 512)), ]) - img = transform(self.Unprompted.shortcode_user_vars["init_images"][0]).unsqueeze(0) + flattened_input = flattened(self.Unprompted.shortcode_user_vars["init_images"][0], opts.img2img_background_color) + img = transform(flattened_input).unsqueeze(0) # predict with torch.no_grad(): @@ -162,4 +165,4 @@ def after(self,p=None,processed=None): processed.images.append(self.image_mask) self.image_mask = None self.show = False - return processed \ No newline at end of file + return processed From 8fbf63b578c5e3782395d992382de1d384231e82 Mon Sep 17 00:00:00 2001 From: Colin Redman <20376935+credman0@users.noreply.github.com> Date: Sun, 25 Dec 2022 22:48:41 -0600 Subject: [PATCH 2/4] Fix typo I made "cleaning it up" on github... I swear I test my code --- shortcodes/stable_diffusion/txt2mask.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shortcodes/stable_diffusion/txt2mask.py b/shortcodes/stable_diffusion/txt2mask.py index 9906e3e..a2b8c67 100644 --- a/shortcodes/stable_diffusion/txt2mask.py +++ b/shortcodes/stable_diffusion/txt2mask.py @@ -101,7 +101,7 @@ def get_mask(): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), transforms.Resize((512, 512)), ]) - flattened_input = flattened(self.Unprompted.shortcode_user_vars["init_images"][0], opts.img2img_background_color) + flattened_input = flatten(self.Unprompted.shortcode_user_vars["init_images"][0], opts.img2img_background_color) img = transform(flattened_input).unsqueeze(0) # predict From 639d1c8157b7f319fce0fe3de97d30a85019840f Mon Sep 17 00:00:00 2001 From: Colin Redman <20376935+credman0@users.noreply.github.com> Date: Mon, 26 Dec 2022 00:21:11 -0600 Subject: [PATCH 3/4] txt2mask pad with cv2 dilation Pad with dilation rather than scaling + cropping --- shortcodes/stable_diffusion/txt2mask.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/shortcodes/stable_diffusion/txt2mask.py b/shortcodes/stable_diffusion/txt2mask.py index a2b8c67..f9953ce 100644 --- a/shortcodes/stable_diffusion/txt2mask.py +++ b/shortcodes/stable_diffusion/txt2mask.py @@ -26,6 +26,14 @@ def run_block(self, pargs, kwargs, context, content): if "smoothing" in kwargs: radius = int(kwargs["smoothing"]) smoothing_kernel = numpy.ones((radius,radius),numpy.float32)/(radius*radius) + + + + # Pad the mask by applying a dilation + mask_padding = int(self.Unprompted.parse_advanced(kwargs["padding"],context) if "padding" in kwargs else 0) + padding_dilation_kernel = None + if (mask_padding > 0): + padding_dilation_kernel = numpy.ones((mask_padding, mask_padding), numpy.uint8) prompts = content.split(self.Unprompted.Config.syntax.delimiter) prompt_parts = len(prompts) @@ -63,6 +71,7 @@ def process_mask_parts(these_preds,these_prompt_parts,mode,final_img = None): img = cv2.imread(filename) if smoothing_kernel is not None: img = cv2.filter2D(img,-1,smoothing_kernel) + if padding_dilation_kernel is not None: img = cv2.dilate(img,padding_dilation_kernel,iterations=1) gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) (thresh, bw_image) = cv2.threshold(gray_image, mask_precision, 255, cv2.THRESH_BINARY) @@ -136,16 +145,7 @@ def get_mask(): if (pixel_data[0] == 0 and pixel_data[1] == 0 and pixel_data[2] == 0): black_pixels += 1 subject_size = 1 - black_pixels / total_pixels self.Unprompted.shortcode_user_vars[kwargs["size_var"]] = subject_size - - # Increase mask size with padding - mask_padding = int(self.Unprompted.parse_advanced(kwargs["padding"],context) if "padding" in kwargs else 0) - if (mask_padding > 0): - aspect_ratio = self.Unprompted.shortcode_user_vars["init_images"][0].width / self.Unprompted.shortcode_user_vars["init_images"][0].height - new_width = self.Unprompted.shortcode_user_vars["init_images"][0].width+mask_padding*2 - new_height = round(new_width / aspect_ratio) - final_img = final_img.resize((new_width,new_height)) - final_img = center_crop(final_img,self.Unprompted.shortcode_user_vars["init_images"][0].width,self.Unprompted.shortcode_user_vars["init_images"][0].height) - + return final_img # Set up processor parameters correctly From 9e3ba8f38210b0f99ae16ad1c7bc65f2c750fbff Mon Sep 17 00:00:00 2001 From: Colin Redman <20376935+credman0@users.noreply.github.com> Date: Mon, 26 Dec 2022 00:32:45 -0600 Subject: [PATCH 4/4] Dilate before smooth --- shortcodes/stable_diffusion/txt2mask.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shortcodes/stable_diffusion/txt2mask.py b/shortcodes/stable_diffusion/txt2mask.py index f9953ce..6fa760a 100644 --- a/shortcodes/stable_diffusion/txt2mask.py +++ b/shortcodes/stable_diffusion/txt2mask.py @@ -70,8 +70,8 @@ def process_mask_parts(these_preds,these_prompt_parts,mode,final_img = None): # TODO: Figure out how to convert the plot above to numpy instead of re-loading image img = cv2.imread(filename) - if smoothing_kernel is not None: img = cv2.filter2D(img,-1,smoothing_kernel) if padding_dilation_kernel is not None: img = cv2.dilate(img,padding_dilation_kernel,iterations=1) + if smoothing_kernel is not None: img = cv2.filter2D(img,-1,smoothing_kernel) gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) (thresh, bw_image) = cv2.threshold(gray_image, mask_precision, 255, cv2.THRESH_BINARY)