new control mask module

CHANGELOG.md

@@ -1,20 +1,21 @@
 # Change Log for SD.Next
 
-## Update for 2023-01-13
+## Update for 2023-01-14
 
 Another release with a lot more functionality in new Control module and FaceID & IPAdapter modules  
 Plus welcome additions to UI performance and accessibility and flexibility of deployment  
 And it also includes fixes for all reported issues so far  
 
-However, 
-
 - **Control**:
   - add **inpaint** support  
     applies to both *img2img* and *controlnet* workflows  
-    *note*: set blur to level you desire  
   - add **outpaint** support  
     applies to both *img2img* and *controlnet* workflows  
     *note*: increase denoising strength since outpainted area is blank by default  
+  - new **mask** module  
+    - granular blur (gaussian), errode (reduce or remove noise) and dilate (pad or expand) with **live preview**  
+    - *optional* **auto-segmentation** (e.g. segment-anything) using ml models  
+      auto segmentation will automatically expand masked area to segments that include current user mask  
   - allow **resize** both *before* and *after* generate operation  
     this allows for workflows such as: *image -> upscale or downscale -> generate -> upscale or downscale -> output*  
     providing more flexibility and than standard hires workflow  

modules/control/run.py

@@ -80,7 +80,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                 hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
                 resize_mode_before, resize_name_before, width_before, height_before, scale_by_before, selected_scale_tab_before,
                 resize_mode_after, resize_name_after, width_after, height_after, scale_by_after, selected_scale_tab_after,
-                denoising_strength, batch_count, batch_size, mask_blur, mask_overlap,
+                denoising_strength, batch_count, batch_size,
                 video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate,
                 ip_adapter, ip_scale, ip_image,
                 *input_script_args
@@ -135,7 +135,6 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
         denoising_strength = denoising_strength,
         n_iter = batch_count,
         batch_size = batch_size,
-        mask_blur=mask_blur,
         outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_control_samples,
         outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_control_grids,
     )

modules/masking.py

@@ -1,4 +1,12 @@
+from types import SimpleNamespace
+import os
+import time
+import gradio as gr
+import numpy as np
+import cv2
 from PIL import Image, ImageFilter, ImageOps
+from transformers import SamModel, SamImageProcessor, MaskGenerationPipeline
+from modules import shared, errors, devices, ui_components, ui_symbols
 
 
 def get_crop_region(mask, pad=0):
@@ -83,3 +91,266 @@ def fill(image, mask):
         for _ in range(repeats):
             image_mod.alpha_composite(blurred)
     return image_mod.convert("RGB")
+
+
+"""
+[docs](https://huggingface.co/docs/transformers/v4.36.1/en/model_doc/sam#overview)
+TODO:
+- PerSAM
+- https://huggingface.co/docs/transformers/tasks/semantic_segmentation
+- transformers.pipeline.MaskGenerationPipeline: https://huggingface.co/models?pipeline_tag=mask-generation
+- transformers.pipeline.ImageSegmentationPipeline: https://huggingface.co/models?pipeline_tag=image-segmentation
+"""
+
+MODELS = {
+    'None': None,
+    'Facebook SAM ViT Base': 'facebook/sam-vit-base',
+    'Facebook SAM ViT Large': 'facebook/sam-vit-large',
+    'Facebook SAM ViT Huge': 'facebook/sam-vit-huge',
+    'SlimSAM Uniform': 'Zigeng/SlimSAM-uniform-50',
+    'SlimSAM Uniform Tiny': 'Zigeng/SlimSAM-uniform-77',
+    # 'Tiny Random': 'fxmarty/sam-vit-tiny-random',
+}
+COLORMAP = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'spring', 'cool', 'hsv', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', 'twilight', 'shifted', 'turbo', 'deepgreen']
+cache_dir = 'models/control/segment'
+loaded_model = None
+model: SamModel = None
+processor: SamImageProcessor = None
+generator: MaskGenerationPipeline = None
+debug = shared.log.trace if os.environ.get('SD_MASK_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: MASK')
+busy = False
+btn_segment = None
+controls = []
+opts = SimpleNamespace(**{
+    'mask_blur': 0.01,
+    'mask_erode': 0.01,
+    'mask_dilate': 0.01,
+    'seg_iou_thresh': 0.5,
+    'seg_score_thresh': 0.5,
+    'seg_nms_thresh': 0.5,
+    'seg_overlap_ratio': 0.3,
+    'seg_points_per_batch': 64,
+    'seg_topK': 50,
+    'seg_colormap': 'pink',
+    'preview_type': 'composite',
+    'seg_live': True,
+    'weight_original': 0.5,
+    'weight_mask': 0.5,
+    'kernel_iterations': 1,
+})
+
+
+def init_model(selected_model: str):
+    global busy, loaded_model, model, processor, generator # pylint: disable=global-statement
+    if selected_model == "None":
+        if model is not None:
+            shared.log.debug('Segment unloading model')
+        model = None
+        loaded_model = None
+        processor = None
+        generator = None
+        devices.torch_gc()
+        return selected_model
+    model_path = MODELS[selected_model]
+    if model_path is not None and (loaded_model != selected_model or model is None or processor is None):
+        busy = True
+        t0 = time.time()
+        shared.log.debug(f'Segment loading: model={selected_model} path={model_path}')
+        model = SamModel.from_pretrained(model_path, cache_dir=cache_dir).to(device=devices.device)
+        processor = SamImageProcessor.from_pretrained(model_path, cache_dir=cache_dir)
+        generator = MaskGenerationPipeline(
+            model=model,
+            image_processor=processor,
+            device=devices.device,
+            # output_bboxes_mask=False,
+            # output_rle_masks=False,
+        )
+        devices.torch_gc()
+        shared.log.debug(f'Segment loaded: model={selected_model} path={model_path} time={time.time()-t0:.2f}s')
+        busy = False
+    return selected_model
+
+
+def run_segment(input_image: gr.Image, input_mask: np.ndarray):
+    outputs = None
+    with devices.inference_context():
+        try:
+            outputs = generator(
+                input_image,
+                points_per_batch=opts.seg_points_per_batch,
+                pred_iou_thresh=opts.seg_iou_thresh,
+                stability_score_thresh=opts.seg_score_thresh,
+                crops_nms_thresh=opts.seg_nms_thresh,
+                crop_overlap_ratio=opts.seg_overlap_ratio,
+                crops_n_layers=0,
+                crop_n_points_downscale_factor=1,
+            )
+        except Exception as e:
+            shared.log.error(f'Segment error: {e}')
+            errors.display(e, 'Segment')
+            return outputs
+    devices.torch_gc()
+    i = 1
+    combined_mask = np.zeros(input_mask.shape, dtype='uint8')
+    input_mask_size = np.count_nonzero(input_mask)
+    debug(f'Segment: {vars(opts)}')
+    for mask in outputs['masks']:
+        mask = mask.astype('uint8')
+        mask_size = np.count_nonzero(mask)
+        if mask_size == 0:
+            continue
+        overlap = 0
+        if input_mask_size > 0:
+            overlap = cv2.bitwise_and(mask, input_mask)
+            overlap = np.count_nonzero(overlap)
+            if overlap == 0:
+                continue
+        mask = (opts.seg_topK + 1 - i) * mask * (255 // opts.seg_topK) # set grayscale intensity so we can recolor
+        combined_mask = combined_mask + mask
+        debug(f'Segment mask: i={i} size={input_image.width}x{input_image.height} masked={mask_size}px overlap={overlap} score={outputs["scores"][i-1]:.2f}')
+        i += 1
+        if i > opts.seg_topK:
+            break
+    return combined_mask
+
+
+def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: str = None, mask_blur: int = None, mask_padding: int = None, segment_enable=True):
+    if input_image is None:
+        return input_mask
+    if isinstance(input_image, list):
+        input_image = input_image[0]
+    if isinstance(input_image, dict):
+        input_mask = input_image.get('mask', None)
+        input_image = input_image.get('image', None)
+    if input_mask is None:
+        input_mask = input_image.convert('L')
+        input_mask = input_mask.point(lambda x: 255 if x > 127 else 0)
+    else:
+        input_mask = input_mask.convert('L')
+    shared.log.debug(f'Segment mask: input={input_image} mask={input_mask} type={return_type}')
+
+    input_mask = np.array(input_mask) // 255
+    t0 = time.time()
+
+    if mask_blur is not None:
+        opts.mask_blur = mask_blur / min(input_image.width, input_image.height)
+    if mask_padding is not None:
+        opts.mask_dilate = mask_padding / min(input_image.width, input_image.height)
+
+    if generator is None or not segment_enable:
+        mask = input_mask * 255
+    else:
+        mask = run_segment(input_image, input_mask)
+
+    if mask is None:
+        shared.log.error('Segment error: no mask')
+        return input_mask
+
+    debug(f'Segment mask: mask={mask.shape}')
+    if opts.mask_erode > 0:
+        try:
+            kernel = np.ones((int(opts.mask_erode * input_image.height / 4) + 1, int(opts.mask_erode * input_image.width / 4) + 1), np.uint8)
+            cv2_mask = cv2.erode(mask, kernel, iterations=opts.kernel_iterations) # remove noise
+            mask = cv2_mask
+            debug(f'Segment erode={opts.mask_erode} kernel={kernel} mask={mask.shape}')
+        except Exception as e:
+            shared.log.error(f'Segment erode: {e}')
+    if opts.mask_dilate > 0:
+        try:
+            kernel = np.ones((int(opts.mask_dilate * input_image.height / 4) + 1, int(opts.mask_dilate * input_image.width / 4) + 1), np.uint8)
+            cv2_mask = cv2.dilate(mask, kernel, iterations=opts.kernel_iterations) # expand area
+            mask = cv2_mask
+            debug(f'Segment dilate={opts.mask_dilate} kernel={kernel} mask={mask.shape}')
+        except Exception as e:
+            shared.log.error(f'Segment dilate: {e}')
+    if opts.mask_blur > 0:
+        try:
+            sigmax, sigmay = 1 + int(opts.mask_blur * input_image.width / 4), 1 + int(opts.mask_blur * input_image.height / 4)
+            cv2_mask = cv2.GaussianBlur(mask, (0, 0), sigmaX=sigmax, sigmaY=sigmay) # blur mask
+            mask = cv2_mask
+            debug(f'Segment blur={opts.mask_blur} x={sigmax} y={sigmay} mask={mask.shape}')
+        except Exception as e:
+            shared.log.error(f'Segment blur: {e}')
+
+    mask_size = np.count_nonzero(mask)
+    total_size = np.prod(mask.shape)
+    area_size = np.count_nonzero(mask)
+    colored_mask = cv2.applyColorMap(mask, COLORMAP.index(opts.seg_colormap)) # recolor mask
+    combined_image = cv2.addWeighted(np.array(input_image), opts.weight_original, colored_mask, opts.weight_mask, 0)
+    binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1] # otsu uses mean instead of threshold
+    t1 = time.time()
+
+    return_type = return_type or opts.preview_type
+    shared.log.debug(f'Segment mask opts: size={input_image.width}x{input_image.height} masked={mask_size}px area={area_size/total_size:.2f} time={t1-t0:.2f}')
+    if return_type == 'none':
+        return input_mask
+    elif return_type == 'binary':
+        return Image.fromarray(binary_mask)
+    elif return_type == 'grayscale':
+        return Image.fromarray(mask)
+    elif return_type == 'color':
+        return Image.fromarray(colored_mask)
+    elif return_type == 'composite':
+        return Image.fromarray(combined_image)
+    return input_mask
+
+
+def run_mask_live(input_image: gr.Image):
+    global busy # pylint: disable=global-statement
+    if opts.seg_live:
+        if not busy:
+            busy = True
+            res = run_mask(input_image)
+            busy = False
+            return res
+    else:
+        return None
+
+
+def create_segment_ui():
+    def update_opts(*args):
+        opts.seg_live = args[0]
+        opts.mask_blur = args[1]
+        opts.mask_erode = args[2]
+        opts.mask_dilate = args[3]
+        opts.seg_score_thresh = args[4]
+        opts.seg_iou_thresh = args[5]
+        opts.seg_nms_thresh = args[6]
+        opts.preview_type = args[7]
+        opts.seg_colormap = args[8]
+
+    def display_controls(selected_model):
+        return 4 * [gr.update(visible=True)] + (len(controls) - 4) * [gr.update(visible=selected_model != 'None')]
+
+    global btn_segment # pylint: disable=global-statement
+    with gr.Accordion(open=False, label="Mask", elem_id="control_mask", elem_classes=["small-accordion"]):
+        controls.clear()
+        with gr.Row():
+            controls.append(gr.Checkbox(label="Live update", value=True))
+        with gr.Row():
+            controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Blur', value=0.01, elem_id="control_mask_blur"))
+            controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Erode', value=0.01, elem_id="control_mask_erode"))
+            controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Dilate', value=0.01, elem_id="control_mask_dilate"))
+        with gr.Row():
+            selected_model = gr.Dropdown(label="Auto-segment", choices=MODELS.keys(), value='None')
+            btn_segment = ui_components.ToolButton(value=ui_symbols.refresh, visible=False)
+        with gr.Row():
+            controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Score', value=0.5, visible=False))
+            controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='IOU', value=0.5, visible=False))
+            controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='NMS', value=0.5, visible=False))
+        with gr.Row():
+            controls.append(gr.Dropdown(label="Preview", choices=['none', 'binary', 'grayscale', 'color', 'composite'], value='composite'))
+            controls.append(gr.Dropdown(label="Colormap", choices=COLORMAP, value='pink'))
+
+        selected_model.change(fn=init_model, inputs=[selected_model], outputs=[selected_model])
+        selected_model.change(fn=display_controls, inputs=[selected_model], outputs=controls)
+        for control in controls:
+            control.change(fn=update_opts, inputs=controls, outputs=[])
+
+
+def bind_controls(input_image: gr.Image, preview_image: gr.Image):
+    btn_segment.click(run_mask, inputs=[input_image], outputs=[preview_image])
+    input_image.edit(fn=run_mask_live, inputs=[input_image], outputs=[preview_image])
+    for control in controls:
+        control.change(fn=run_mask_live, inputs=[input_image], outputs=[preview_image])

modules/processing.py

@@ -1290,15 +1290,18 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         if self.image_mask is not None:
             if type(self.image_mask) == list:
                 self.image_mask = self.image_mask[0]
-            self.image_mask = create_binary_mask(self.image_mask)
-            if self.inpainting_mask_invert:
-                self.image_mask = ImageOps.invert(self.image_mask)
-            if self.mask_blur > 0:
-                np_mask = np.array(self.image_mask)
-                kernel_size = 2 * int(2.5 * self.mask_blur + 0.5) + 1
-                np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur)
-                np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur)
-                self.image_mask = Image.fromarray(np_mask)
+            if shared.backend == shared.Backend.ORIGINAL:
+                self.image_mask = create_binary_mask(self.image_mask)
+                if self.inpainting_mask_invert:
+                    self.image_mask = ImageOps.invert(self.image_mask)
+                if self.mask_blur > 0:
+                    np_mask = np.array(self.image_mask)
+                    kernel_size = 2 * int(2.5 * self.mask_blur + 0.5) + 1
+                    np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur)
+                    np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur)
+                    self.image_mask = Image.fromarray(np_mask)
+            else:
+                self.image_mask = modules.masking.run_mask(input_image=self.init_images, input_mask=self.image_mask, return_type='grayscale', mask_blur=self.mask_blur, mask_padding=self.inpaint_full_res_padding, segment_enable=False)
             if self.inpaint_full_res:
                 self.mask_for_overlay = self.image_mask
                 mask = self.image_mask.convert('L')

modules/processing_diffusers.py

@@ -17,6 +17,7 @@ import modules.prompt_parser_diffusers as prompt_parser_diffusers
 from modules.sd_hijack_hypertile import hypertile_set
 from modules.processing_correction import correction_callback
 from modules.processing_vae import vae_encode, vae_decode
+from modules.masking import run_mask
 
 
 debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -167,9 +168,10 @@ def process_diffusers(p: StableDiffusionProcessing):
         elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images' ,[])) > 0:
             p.ops.append('inpaint')
             if p.task_args.get('mask_image', None) is not None: # provided as override by a module
-                p.mask = shared.sd_model.mask_processor.blur(p.task_args['mask_image'], blur_factor=p.mask_blur) if p.mask_blur > 0 else p.task_args['mask_image']
+                # p.mask = shared.sd_model.mask_processor.blur(p.task_args['mask_image'], blur_factor=p.mask_blur) if p.mask_blur > 0 else p.task_args['mask_image']
+                p.mask = run_mask(input_image=p.init_images, input_mask=p.task_args['mask_image'], return_type='grayscale')
             elif getattr(p, 'image_mask', None) is not None: # standard
-                p.mask = p.image_mask
+                p.mask = run_mask(input_image=p.init_images, input_mask=p.image_mask, return_type='grayscale')
             elif getattr(p, 'mask', None) is not None: # backward compatibility
                 pass
             else: # fallback

modules/segment.py

@@ -1,157 +0,0 @@
-"""
-[docs](https://huggingface.co/docs/transformers/v4.36.1/en/model_doc/sam#overview)
-TODO:
-- PerSAM
-- transformers.pipeline.MaskGenerationPipeline: https://huggingface.co/models?pipeline_tag=mask-generation
-- transformers.pipeline.ImageSegmentationPipeline: https://huggingface.co/models?pipeline_tag=image-segmentation
-"""
-
-from transformers import SamModel, SamImageProcessor, MaskGenerationPipeline
-from PIL import Image
-import gradio as gr
-import numpy as np
-import cv2
-from modules import shared, devices
-
-
-MODELS = {
-    'None': None,
-    'Facebook SAM ViT Base': 'facebook/sam-vit-base',
-    'Facebook SAM ViT Large': 'facebook/sam-vit-large',
-    'Facebook SAM ViT Huge': 'facebook/sam-vit-huge',
-    'SlimSAM Uniform': 'Zigeng/SlimSAM-uniform-50',
-}
-COLORMAP = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'spring', 'cool', 'hsv', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', 'twilight', 'shifted', 'turbo', 'deepgreen']
-cache_dir = 'models/control/segment'
-loaded_model = None
-model: SamModel = None
-processor: SamImageProcessor = None
-
-
-def init(selected_model: str, input_image: gr.Image):
-    global loaded_model, model, processor # pylint: disable=global-statement
-    if input_image is None or input_image.get('image', None) is None:
-        return False
-    if selected_model == "None":
-        return False
-    if selected_model == "None":
-        model = None
-        loaded_model = None
-        processor = None
-    model_path = MODELS[selected_model]
-    if model_path is not None and (loaded_model != selected_model or model is None or processor is None):
-        shared.log.debug(f'Segment loading: model={selected_model} path={model_path}')
-        model = SamModel.from_pretrained(model_path, cache_dir=cache_dir).to(device=devices.device)
-        processor = SamImageProcessor.from_pretrained(model_path, cache_dir=cache_dir)
-        shared.log.debug(f'Segment loaded: model={selected_model} path={model_path}')
-    if model is None or processor is None:
-        return False
-    return True
-
-
-# run as auto-mask with all possible masks
-def run_segment(selected_model: str, input_image: gr.Image, points_per_batch=64, pred_iou_thresh=0.75, stability_score_thresh=0.85, crops_nms_thresh=0.5, crop_overlap_ratio=0.3, topK=25, colormap='jet', erode=0, dilate=0):
-    if not init(selected_model, input_image):
-        return gr.update(), None
-    input_mask = input_image.get('mask', None) or Image.new('L', input_image.get('image', None).size, 255)
-    input_mask = input_mask.convert('L')
-    input_image = input_image.get('image', None)
-    generator: MaskGenerationPipeline = MaskGenerationPipeline(model=model, image_processor=processor, device=devices.device)
-    with devices.inference_context():
-        outputs = generator(
-            input_image,
-            points_per_batch=points_per_batch,
-            pred_iou_thresh=pred_iou_thresh,
-            stability_score_thresh=stability_score_thresh,
-            crops_nms_thresh=crops_nms_thresh,
-            crop_overlap_ratio=crop_overlap_ratio,
-        )
-    combined_mask = np.zeros(input_mask.size, dtype='uint8')
-    input_mask = np.array(input_mask) // 255
-    input_mask_size = np.count_nonzero(input_mask)
-    print('HERE', input_mask.shape, input_mask_size)
-    i = 1
-    for mask in outputs['masks']:
-        mask = mask.astype('uint8')
-        mask_size = np.count_nonzero(mask)
-        if mask_size == 0:
-            continue
-        overlap = 0
-        if input_mask_size > 0:
-            overlap = cv2.bitwise_and(mask, input_mask)
-            overlap = np.count_nonzero(overlap)
-            if overlap == 0:
-                continue
-        # TODO erode,dilate
-        if erode > 0:
-            mask = cv2.erode(mask, np.ones((erode, erode), np.uint8), iterations=2) # remove noise
-        if dilate > 0:
-            mask = cv2.dilate(mask, np.ones((dilate, dilate), np.uint8), iterations=2) # expand area
-        mask = (topK + 1 - i) * mask * (255 // topK) # set grayscale intensity so we can recolor
-        combined_mask = combined_mask + mask
-        i += 1
-        if i > topK:
-            break
-    mask_size = np.count_nonzero(combined_mask)
-    total_size = np.prod(combined_mask.shape)
-    area_size = np.count_nonzero(combined_mask)
-    shared.log.debug(f'Segment mask: size={input_image.width}x{input_image.height} input={input_mask_size}px masked={mask_size}px area={area_size/total_size:.2f}')
-    colored_mask = cv2.applyColorMap(combined_mask, COLORMAP.index(colormap)) # recolor mask
-    combined_image = cv2.addWeighted(np.array(input_image), 0.6, colored_mask, 0.4, 0)
-    _thres, binary_mask = cv2.threshold(combined_mask, 1, 255, cv2.THRESH_BINARY_INV) # create mask
-    binary_mask = np.invert(binary_mask)
-
-    binary_mask = Image.fromarray(binary_mask)
-    combined_mask = Image.fromarray(combined_mask)
-    colored_mask = Image.fromarray(colored_mask)
-    overlay_image = Image.fromarray(combined_image)
-
-    # TODO return type
-    binary_mask.save('/tmp/mask-binary.png')
-    combined_mask.save('/tmp/mask-combined.png')
-    combined_mask.save('/tmp/mask-colored.png')
-    overlay_image.save('/tmp/mask-overlay.png')
-    return input_image, overlay_image
-
-
-# run with sam model directly needing set of points
-def run_segment_points(selected_model: str, input_image: gr.Image):
-    if not init(selected_model, input_image):
-        return input_image
-    # input_mask = input_image.get('mask', None) or Image.new('L', input_image.get('image', None).size, 0)
-    input_image = input_image.get('image', None)
-    with devices.inference_context():
-        inputs = processor(
-            input_image,
-            input_points=[[[256, 256]]], # TODO calculate points based on mask
-            return_tensors="pt"
-        ).to(device=devices.device)
-        outputs = model(
-            pixel_values=inputs['pixel_values'],
-            multimask_output=True,
-            )
-        masks = processor.post_process_masks(
-            outputs.pred_masks.cpu(),
-            inputs["original_sizes"].cpu(),
-            inputs["reshaped_input_sizes"].cpu()
-        )
-        scores = outputs.iou_scores
-    mask = masks[0].squeeze(0)
-    scores = scores[0].squeeze(0)
-    masks = mask.unbind(0)
-    output_masks = []
-    for i, mask in enumerate(masks):
-        mask = mask.detach().cpu().numpy()
-        mask = mask.astype('uint8') * 255
-        mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2)
-        total_size = np.prod(mask.shape)
-        area_size = np.count_nonzero(mask)
-        shared.log.debug(f'Segment mask: area={area_size/total_size:.2f} score={scores[i].item():.2f}')
-        mask = Image.fromarray(mask)
-        output_masks.append(mask)
-
-
-def create_segment_ui(input_image: gr.Image, preview_image: gr.Image):
-    selected = gr.Dropdown(label="Segment", choices=MODELS.keys(), value='None')
-    selected.change(fn=run_segment, inputs=[selected, input_image], outputs=[input_image, preview_image])
-    return selected

modules/ui_control.py

@@ -2,7 +2,6 @@ import os
 import time
 import gradio as gr
 import matplotlib.pyplot as plt
-import numpy as np
 from PIL import Image
 from modules.control import unit
 from modules.control import processors # patrickvonplaten controlnet_aux
@@ -12,7 +11,7 @@ from modules.control.units import lite # vislearn ControlNet-XS
 from modules.control.units import t2iadapter # TencentARC T2I-Adapter
 from modules.control.units import reference # reference pipeline
 from scripts import ipadapter # pylint: disable=no-name-in-module
-from modules import errors, shared, progress, sd_samplers, ui_components, ui_symbols, ui_common, ui_sections, generation_parameters_copypaste, call_queue, scripts, segment # pylint: disable=ungrouped-imports
+from modules import errors, shared, progress, sd_samplers, ui_components, ui_symbols, ui_common, ui_sections, generation_parameters_copypaste, call_queue, scripts, masking # pylint: disable=ungrouped-imports
 
 
 gr_height = 512
@@ -34,6 +33,7 @@ def initialize():
     lite.cache_dir = os.path.join(shared.opts.control_dir, 'lite')
     t2iadapter.cache_dir = os.path.join(shared.opts.control_dir, 'adapter')
     processors.cache_dir = os.path.join(shared.opts.control_dir, 'processor')
+    masking.cache_dir   = os.path.join(shared.opts.control_dir, 'segment')
     unit.default_device = devices.device
     unit.default_dtype = devices.dtype
     os.makedirs(shared.opts.control_dir, exist_ok=True)
@@ -42,6 +42,7 @@ def initialize():
     os.makedirs(lite.cache_dir, exist_ok=True)
     os.makedirs(t2iadapter.cache_dir, exist_ok=True)
     os.makedirs(processors.cache_dir, exist_ok=True)
+    os.makedirs(masking.cache_dir, exist_ok=True)
     scripts.scripts_current = scripts.scripts_control
     scripts.scripts_current.initialize_scripts(is_img2img=True)
 
@@ -116,35 +117,7 @@ def get_video(filepath: str):
         return msg
 
 
-def select_mask(image: Image.Image, negative: bool = False):
-    if image is None:
-        return image
-    image_mask = image.convert("L")
-    if negative:
-        image_mask = image_mask.point(lambda x: 255 if x < 4 else 0)
-    else:
-        image_mask = image_mask.point(lambda x: 255 if x > 127 else 0)
-    return image_mask
-
-
-def expand_mask(image: Image.Image, expand: int = 64):
-    import cv2
-    if image is None:
-        return image
-    pil_mask = image.convert("L")
-    np_mask = np.array(pil_mask)
-    erode, dilate, threshold = 3, 8, 4
-    if threshold > 0:
-        _thres, np_mask = cv2.threshold(np_mask, threshold, 255, cv2.THRESH_BINARY_INV) # create mask
-    if erode > 0:
-        np_mask = cv2.erode(np_mask, np.ones((erode, erode), np.uint8), iterations=expand//dilate) # remove noise
-    if dilate > 0:
-        np_mask = cv2.dilate(np_mask, np.ones((dilate, dilate), np.uint8), iterations=expand//dilate) # expand area
-    image_mask = Image.fromarray(np_mask.astype(np.uint8))
-    return image_mask
-
-
-def select_input(input_mode, input_image, selected_init, init_type, input_resize, input_inpaint, input_video, input_batch, input_folder, _mask_blur, mask_overlap):
+def select_input(input_mode, input_image, selected_init, init_type, input_resize, input_inpaint, input_video, input_batch, input_folder):
     global busy, input_source, input_init, input_mask # pylint: disable=global-statement
     busy = True
     if input_mode == 'Select':
@@ -172,14 +145,14 @@ def select_input(input_mode, input_image, selected_init, init_type, input_resize
     # control inputs
     if isinstance(selected_input, Image.Image): # image via upload -> image
         if input_mode == 'Outpaint':
-            input_mask = expand_mask(image=selected_input, expand=mask_overlap)
+            input_mask = masking.run_mask(input_image=selected_input, input_mask=None, return_type='grayscale')
         input_source = [selected_input]
         input_type = 'PIL.Image'
         shared.log.debug(f'Control input: type={input_type} input={input_source}')
         status = f'Control input | Image | Size {selected_input.width}x{selected_input.height} | Mode {selected_input.mode}'
         res = [gr.Tabs.update(selected='out-gallery'), status]
     elif isinstance(selected_input, dict): # inpaint -> dict image+mask
-        input_mask = select_mask(image=selected_input['mask'], negative=False)
+        input_mask = masking.run_mask(input_image=selected_input['image'], input_mask=selected_input['mask'], return_type='grayscale')
         selected_input = selected_input['image']
         input_source = [selected_input]
         input_type = 'PIL.Image'
@@ -217,7 +190,7 @@ def select_input(input_mode, input_image, selected_init, init_type, input_resize
     elif init_type == 2: # Separate init image
         if isinstance(selected_init, Image.Image): # image via upload -> image
             if input_mode == 'Outpaint':
-                input_mask = expand_mask(image=selected_init, expand=mask_overlap)
+                input_mask = masking.run_mask(input_image=selected_init, input_mask=None, return_type='grayscale')
             input_source = [selected_init]
             input_init = [selected_init]
             input_type = 'PIL.Image'
@@ -225,7 +198,7 @@ def select_input(input_mode, input_image, selected_init, init_type, input_resize
             status = f'Control input | Image | Size {selected_init.width}x{selected_init.height} | Mode {selected_init.mode}'
             res = [gr.Tabs.update(selected='out-gallery'), status]
         elif isinstance(selected_init, dict): # inpaint -> dict image+mask
-            input_mask = select_mask(image=selected_init['mask'])
+            input_mask = masking.run_mask(input_image=selected_init['image'], input_mask=selected_init['mask'], return_type='grayscale')
             input_init = selected_init['image']
             input_source = [selected_init]
             input_type = 'PIL.Image'
@@ -313,9 +286,6 @@ def create_ui(_blocks: gr.Blocks=None):
                         input_type = gr.Radio(label="Input type", choices=['Control only', 'Init image same as control', 'Separate init image'], value='Control only', type='index', elem_id='control_input_type')
                     with gr.Row():
                         denoising_strength = gr.Slider(minimum=0.01, maximum=1.0, step=0.01, label='Denoising strength', value=0.50, elem_id="control_denoising_strength")
-                    with gr.Row():
-                        mask_blur = gr.Slider(minimum=0, maximum=100, step=1, label='Blur', value=8, elem_id="control_mask_blur")
-                        mask_overlap = gr.Slider(minimum=0, maximum=100, step=1, label='Overlap', value=64, elem_id="control_mask_overlap")
 
                 with gr.Accordion(open=False, label="Size", elem_id="control_size", elem_classes=["small-accordion"]):
                     with gr.Tabs():
@@ -329,7 +299,11 @@ def create_ui(_blocks: gr.Blocks=None):
                     steps, sampler_index = ui_sections.create_sampler_and_steps_selection(sd_samplers.samplers, "control")
 
                 batch_count, batch_size = ui_sections.create_batch_inputs('control')
+
                 seed, _reuse_seed, subseed, _reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = ui_sections.create_seed_inputs('control', reuse_visible=False)
+
+                masking.create_segment_ui()
+
                 cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('control')
 
                 with gr.Accordion(open=False, label="Video", elem_id="control_video", elem_classes=["small-accordion"]):
@@ -366,7 +340,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             input_mode = gr.Label(value='select', visible=False)
                             input_image = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="editor", height=gr_height, visible=True, image_mode='RGB', elem_id='control_input_select')
                             input_resize = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="select", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_resize')
-                            input_inpaint = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="sketch", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_inpaint', brush_radius=64, mask_opacity=0.6)
+                            input_inpaint = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="sketch", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_inpaint', brush_radius=32, mask_opacity=0.6)
                             interrogate_clip, interrogate_booru = ui_sections.create_interrogate_buttons('control')
                             with gr.Row():
                                 input_buttons = [gr.Button('Select', visible=True, interactive=False), gr.Button('Inpaint', visible=True, interactive=True), gr.Button('Outpaint', visible=True, interactive=True)]
@@ -402,10 +376,6 @@ def create_ui(_blocks: gr.Blocks=None):
                         with gr.Tab('Preview', id='preview-image') as tab_image:
                             preview_process = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=gr_height, visible=True)
 
-            # TODO segment as accordian
-            # with gr.Row():
-            #    segment_ui = segment.create_segment_ui(input_inpaint, preview_process)
-
             with gr.Tabs(elem_id='control-tabs') as _tabs_control_type:
 
                 with gr.Tab('ControlNet') as _tab_controlnet:
@@ -684,7 +654,7 @@ def create_ui(_blocks: gr.Blocks=None):
                 interrogate_clip.click(fn=ui_common.interrogate_clip, inputs=[input_image], outputs=[prompt])
                 interrogate_booru.click(fn=ui_common.interrogate_booru, inputs=[input_image], outputs=[prompt])
 
-                select_fields = [input_mode, input_image, init_image, input_type, input_resize, input_inpaint, input_video, input_batch, input_folder, mask_blur, mask_overlap]
+                select_fields = [input_mode, input_image, init_image, input_type, input_resize, input_inpaint, input_video, input_batch, input_folder]
                 select_output = [output_tabs, result_txt]
                 select_dict = dict(
                     fn=select_input,
@@ -711,7 +681,7 @@ def create_ui(_blocks: gr.Blocks=None):
                     cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
                     resize_mode_before, resize_name_before, width_before, height_before, scale_by_before, selected_scale_tab_before,
                     resize_mode_after, resize_name_after, width_after, height_after, scale_by_after, selected_scale_tab_after,
-                    denoising_strength, batch_count, batch_size, mask_blur, mask_overlap,
+                    denoising_strength, batch_count, batch_size,
                     video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate,
                     ip_adapter, ip_scale, ip_image,
                 ]
@@ -736,6 +706,8 @@ def create_ui(_blocks: gr.Blocks=None):
                 generation_parameters_copypaste.add_paste_fields("control", input_image, paste_fields, override_settings)
                 bindings = generation_parameters_copypaste.ParamBinding(paste_button=btn_paste, tabname="control", source_text_component=prompt, source_image_component=output_gallery)
                 generation_parameters_copypaste.register_paste_params_button(bindings)
+                masking.bind_controls(input_inpaint, preview_process)
+
 
                 if os.environ.get('SD_CONTROL_DEBUG', None) is not None: # debug only
                     from modules.control.test import test_processors, test_controlnets, test_adapters, test_xs, test_lite