In img2img batch mode no mask create

[ShmuelRonen]

I want to upload a folder of numbered PNG files and receive the masks in binary, when they are numbered. Is it possible?

[robbieprevost]

I was trying to do the same thing, but looking at the code. It's doesn't look like there's anything in there to save the masks. I'm going to try to edit the python file to see if I can get it to save every mask created, but I'm not great with python so it might be messy. I'll come back and let you know if I can get it to work.

[robbieprevost]

I got it to save the masks after some trial and error by changing bg_mask.py to this:

import gradio as gr
import rembg
from modules import scripts
from modules.ui_components import InputAccordion, FormRow
from modules.processing import StableDiffusionProcessing
import imageio
from datetime import datetime
import os
import re
import numpy as np
from pathlib import Path

models = [
    "None",
    "u2net",
    "u2netp",
    "u2net_human_seg",
    "u2net_cloth_seg",
    "silueta",
    "isnet-general-use",
    "isnet-anime",
]


class BackgroundMaskScript(scripts.Script):
    # 对图片进行放大扩充的脚本
    id_part = "bg_mask"

    def __init__(self):
        super().__init__()
        self.mask = None

    def title(self):
        return "Background Mask"

    def show(self, is_img2img):
        if is_img2img:
            return scripts.AlwaysVisible
        else:
            return False

    def ui(self, is_img2img):
        id_part = self.id_part
        with gr.Accordion(label=self.title(), elem_id=f"{id_part}_main_accordion",open=False):
            with FormRow():
                model = gr.Dropdown(choices=models, value="u2net", show_label=False)
                enable_i2irembg = gr.Checkbox(label="Enable", value=False)
                return_mask = gr.Checkbox(label="Return mask", value=False)
                alpha_matting = gr.Checkbox(label="Alpha matting", value=False)

            with FormRow(visible=False) as alpha_mask_row:
                alpha_matting_erode_size = gr.Slider(label="Erode size", minimum=0, maximum=40, step=1, value=10)
                alpha_matting_foreground_threshold = gr.Slider(label="Foreground threshold", minimum=0, maximum=255,
                                                               step=1, value=240)
                alpha_matting_background_threshold = gr.Slider(label="Background threshold", minimum=0, maximum=255,
                                                               step=1, value=10)

            alpha_matting.change(
                fn=lambda x: gr.update(visible=x),
                inputs=[alpha_matting],
                outputs=[alpha_mask_row],
            )
        return [enable_i2irembg, model, return_mask, alpha_matting, alpha_matting_erode_size,
                alpha_matting_foreground_threshold, alpha_matting_background_threshold]

    def before_process(self, p, *args):
        """
        预处理
        """
        enable_i2irembg, model, return_mask, alpha_matting, alpha_matting_erode_size, alpha_matting_foreground_threshold, alpha_matting_background_threshold, *more = args
        if not enable_i2irembg:
            return

        if not model or model == "None":
            return

        image = p.init_images[0]
        # image = images.resize_image(p.resize_mode, image, p.width, p.height)

        mask = rembg.remove(
            image,
            session=rembg.new_session(model),
            only_mask=True,
            alpha_matting=alpha_matting,
            alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
            alpha_matting_background_threshold=alpha_matting_background_threshold,
            alpha_matting_erode_size=alpha_matting_erode_size,
        )
        self.mask = mask

        p.image_mask = mask
        p.mask_blur_x = 1
        p.mask_blur_y = 1
        p.inpaint_full_res_padding = 32  # 预留像素

        p.inpainting_mask_invert = 0  # 蒙版模式 0=重绘蒙版内容
        p.inpainting_fill = 1  # 蒙版蒙住的内容 0=填充
        p.inpaint_full_res = 0  # 重绘区域 0 = 全图 1=仅蒙版区域

    def find_highest_file_number(self, directory, prefix):
        pattern = re.compile(f"{prefix}_(\d+)\\.png")

        existing_numbers = [int(pattern.match(filename).group(1)) for filename in os.listdir(directory) if pattern.match(filename)]
        return max(existing_numbers, default=0)

    def postprocess(self, p: StableDiffusionProcessing, processed, *args):
        enable_i2irembg, model, return_mask, *more = args

        if not enable_i2irembg:
            return

        if not model or model == "None":
            return

        if not return_mask:
            return
        date = datetime.now().strftime("%Y-%m-%d")
        output_dir = Path(f"{p.outpath_samples}/BGmask/{date}")
        output_dir.mkdir(parents=True, exist_ok=True)
        output_prefix = "mask_output"
        counter = self.find_highest_file_number(output_dir, output_prefix) + 1
        output_path = os.path.join(output_dir, f"{output_prefix}_{counter}.png")
        
        imageio.imwrite(output_path, np.array(self.mask))

        processed.images.append(self.mask)

I can't guarantee that it'll work in your setup, but I figured I'd share it because it's working for me.

[ShmuelRonen]

Thanks! I will check it out.