diff --git a/modules/masking.py b/modules/masking.py index c31f3d27d..fac742d78 100644 --- a/modules/masking.py +++ b/modules/masking.py @@ -135,7 +135,6 @@ opts = SimpleNamespace(**{ 'auto_mask': 'None', 'mask_only': False, 'mask_blur': 0.01, - 'mask_padding': 0, 'mask_erode': 0.01, 'mask_dilate': 0.01, 'seg_iou_thresh': 0.5, @@ -159,7 +158,7 @@ def init_model(selected_model: str): model_path = MODELS[selected_model] if model_path is None: # none if generator is not None: - shared.log.debug('Segment unloading model') + shared.log.debug('Mask segment unloading model') opts.model = None generator = None devices.torch_gc() @@ -172,7 +171,7 @@ def init_model(selected_model: str): if opts.model != selected_model or generator is None: # sam pipeline busy = True t0 = time.time() - shared.log.debug(f'Segment loading: model={selected_model} path={model_path}') + shared.log.debug(f'Mask 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( @@ -183,7 +182,7 @@ def init_model(selected_model: str): # 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') + shared.log.debug(f'Mask segment loaded: model={selected_model} path={model_path} time={time.time()-t0:.2f}s') opts.model = selected_model busy = False return selected_model @@ -204,8 +203,8 @@ def run_segment(input_image: gr.Image, input_mask: np.ndarray): crop_n_points_downscale_factor=1, ) except Exception as e: - shared.log.error(f'Segment error: {e}') - errors.display(e, 'Segment') + shared.log.error(f'Mask segment error: {e}') + errors.display(e, 'Mask segment') return outputs devices.torch_gc() i = 1 @@ -238,7 +237,7 @@ def run_rembg(input_image: Image, input_mask: np.ndarray): try: import rembg except Exception as e: - shared.log.error(f'Segment Rembg load failed: {e}') + shared.log.error(f'Mask Rembg load failed: {e}') return input_mask if "U2NET_HOME" not in os.environ: os.environ["U2NET_HOME"] = os.path.join(paths.models_path, "Rembg") @@ -334,15 +333,14 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st if input_mask is None: return None + size = min(input_image.width, input_image.height) + debug(f'Mask args: blur={mask_blur} padding={mask_padding}') if invert is not None: opts.invert = invert - if mask_blur is not None: # compatibility with old img2img values which have different range - 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) - opts.mask_padding = mask_padding - else: - opts.mask_padding = int(opts.mask_dilate * input_image.height / 4) + 1 + if mask_blur is not None: # compatibility with old img2img values which uses px values + opts.mask_blur = round(4 * mask_blur / size, 3) + if mask_padding is not None: # compatibility with old img2img values which uses px values + opts.mask_erode = 4 * mask_padding / size if opts.model is None or not segment_enable: mask = input_mask @@ -352,32 +350,31 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st mask = run_segment(input_image, input_mask) mask = cv2.resize(mask, (input_image.width, input_image.height), interpolation=cv2.INTER_LINEAR) - debug(f'Mask opts: {opts}') - debug(f'Segment mask: mask={mask.shape}') + debug(f'Mask shape={mask.shape} opts={opts}') 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) + kernel = np.ones((int(opts.mask_erode * size / 4) + 1, int(opts.mask_erode * size / 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.shape} mask={mask.shape}') + debug(f'Mask erode={opts.mask_erode} kernel={kernel.shape} mask={mask.shape}') except Exception as e: - shared.log.error(f'Segment erode: {e}') + shared.log.error(f'Mask 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) + kernel = np.ones((int(opts.mask_dilate * size / 4) + 1, int(opts.mask_dilate * size / 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.shape} mask={mask.shape}') + debug(f'Mask dilate={opts.mask_dilate} kernel={kernel.shape} mask={mask.shape}') except Exception as e: - shared.log.error(f'Segment dilate: {e}') + shared.log.error(f'Mask 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) + sigmax, sigmay = 1 + int(opts.mask_blur * size / 4), 1 + int(opts.mask_blur * size / 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}') + debug(f'Mask blur={opts.mask_blur} x={sigmax} y={sigmay} mask={mask.shape}') except Exception as e: - shared.log.error(f'Segment blur: {e}') + shared.log.error(f'Mask blur: {e}') if opts.invert: mask = np.invert(mask) @@ -388,7 +385,7 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st return_type = return_type or opts.preview_type - shared.log.debug(f'Segment mask: size={input_image.width}x{input_image.height} masked={mask_size}px area={area_size/total_size:.2f} auto={opts.auto_mask} type={return_type} time={t1-t0:.2f}') + shared.log.debug(f'Mask: size={input_image.width}x{input_image.height} masked={mask_size}px area={area_size/total_size:.2f} auto={opts.auto_mask} type={return_type} time={t1-t0:.2f}') if return_type == 'None': return input_mask elif return_type == 'Binary': @@ -410,7 +407,7 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st combined_image = cv2.addWeighted(orig, opts.weight_original, colored_mask, opts.weight_mask, 0) return Image.fromarray(combined_image) else: - shared.log.error(f'Segment unknown return type: {return_type}') + shared.log.error(f'Mask unknown return type: {return_type}') return input_mask @@ -424,9 +421,9 @@ def run_lama(input_image: gr.Image, input_mask: gr.Image = None): input_mask = run_mask(input_image, input_mask, return_type='Grayscale') if lama_model is None: import modules.lama - shared.log.debug(f'LaMa loading: model={modules.lama.LAMA_MODEL_URL}') + shared.log.debug(f'Mask LaMa loading: model={modules.lama.LAMA_MODEL_URL}') lama_model = modules.lama.SimpleLama() - shared.log.debug(f'LaMa loaded: {memory_stats()}') + shared.log.debug(f'Mask LaMa loaded: {memory_stats()}') lama_model.model.to(devices.device) result = lama_model(input_image, input_mask) if shared.opts.control_move_processor: