update interface

.gitignore

@@ -1,3 +1,3 @@
 .idea/
 __pycache__/
-playground.py
+playground.py

app.py

@@ -1,8 +1,9 @@
-import PIL.Image
+import logging
+
 import gradio as gr
-from PIL import ImageColor
-from rembg import new_session, remove
+from rembg import new_session
 
+from cutter import remove, make_label
 from utils import *
 
 remove_bg_models = {
@@ -14,40 +15,53 @@ remove_bg_models = {
 model_choices = keys(remove_bg_models)
 
 
-def alpha_matting(state):
-    if state:
-        return 270, 20, 11
-    return 0, 0, 0
-
+def predict(image, session, smoot, matting, bg_color):
 
-def predict(image,
-            session,
-            matting,
-            only_mask,
-            post_process_mask,
-            foreground_threshold,
-            background_threshold,
-            matting_erode_size,
-            new_bg_color,
-            bg_color,
-            transparency):
     session = new_session(remove_bg_models[session])
 
-    if new_bg_color:
-        r, g, b, _ = ImageColor.getcolor(bg_color, "RGBA")
-        bg_color = r, g, b, transparency
-    else:
-        bg_color = None
+    try:
+        return remove(session, image, smoot, matting, bg_color)
+    except ValueError as err:
+        logging.error(err)
+        return make_label(str(err)), None
+
+
+def change_show_mask(chk_state):
+    return gr.Image.update(visible=chk_state)
+
+
+def change_include_matting(chk_state):
+    return gr.Box.update(visible=chk_state), (0, 0, 0), 0, 0, 0
+
+
+def change_foreground_threshold(fg_value, value):
+    fg, bg, erode = value
+    return fg_value, bg, erode
+
+
+def change_background_threshold(bg_value, value):
+    fg, bg, erode = value
+    return fg, bg_value, erode
+
+
+def change_erode_size(erode_value, value):
+    fg, bg, erode = value
+    return fg, bg, erode_value
+
 
-    return remove(data=image,
-                  session=session,
-                  alpha_matting=matting,
-                  only_mask=only_mask,
-                  post_process_mask=post_process_mask,
-                  bgcolor=bg_color,
-                  alpha_matting_foreground_threshold=foreground_threshold,
-                  alpha_matting_background_threshold=background_threshold,
-                  alpha_matting_erode_size=matting_erode_size)
+def set_dominant_color(chk_state):
+    return chk_state, gr.ColorPicker.update(value=False, visible=not chk_state)
+
+
+def change_picker_color(picker, dominant):
+    if not dominant:
+        return picker
+    return dominant
+
+
+def change_background_mode(chk_state):
+    return gr.ColorPicker.update(value=False, visible=chk_state), \
+        gr.Checkbox.update(value=False, visible=chk_state)
 
 
 footer = r"""
@@ -59,33 +73,57 @@ Demo based on <a href='https://github.com/danielgatis/rembg'>Rembg</a>
 """
 
 with gr.Blocks(title="Remove background") as app:
+    color_state = gr.State(value=False)
+    matting_state = gr.State(value=(0, 0, 0))
+
     gr.HTML("<center><h1>Remove Background Tool</h1></center>")
     with gr.Row(equal_height=False):
         with gr.Column():
-            input_img = gr.Image(type="numpy", label="Input image")
+            input_img = gr.Image(type="pil", label="Input image")
             drp_models = gr.Dropdown(choices=model_choices, label="Model Segment", value="U2NET")
             with gr.Row():
-                chk_alm = gr.Checkbox(label="Alpha Matting", value=False)
-                chk_psm = gr.Checkbox(label="Post process mask", value=False)
-                chk_msk = gr.Checkbox(label="Only Mask", value=False)
-            sld_aft = gr.Slider(0, 300, value=0, step=1, label="Alpha matting foreground threshold")
-            sld_amb = gr.Slider(0, 50, value=0, step=1, label="Alpha matting background threshold")
-            sld_aes = gr.Slider(0, 20, value=0, step=1, label="Alpha matting erode size")
+                chk_include_matting = gr.Checkbox(label="Matting", value=False)
+                chk_smoot_mask = gr.Checkbox(label="Smoot Mask", value=False)
+                chk_show_mask = gr.Checkbox(label="Show Mask", value=False)
+            with gr.Box(visible=False) as slider_matting:
+                slr_fg_threshold = gr.Slider(0, 300, value=270, step=1, label="Alpha matting foreground threshold")
+                slr_bg_threshold = gr.Slider(0, 50, value=20, step=1, label="Alpha matting background threshold")
+                slr_erode_size = gr.Slider(0, 20, value=11, step=1, label="Alpha matting erode size")
             with gr.Box():
                 with gr.Row():
-                    chk_col = gr.Checkbox(label="Change background color", value=False)
-                    color = gr.ColorPicker(label="Pick a new color")
-                    trans = gr.Number(label="Transparency level", value=255, precision=0, minimum=0, maximum=255)
+                    chk_change_color = gr.Checkbox(label="Change background color", value=False)
+                    pkr_color = gr.ColorPicker(label="Pick a new color", visible=False)
+                    chk_dominant = gr.Checkbox(label="Use dominant color", value=False, visible=False)
             run_btn = gr.Button(value="Remove background", variant="primary")
         with gr.Column():
-            output_img = gr.Image(type="pil", label="result")
-            gr.ClearButton(components=[input_img, output_img])
+            output_img = gr.Image(type="pil", label="Image Result")
+            mask_img = gr.Image(type="pil", label="Image Mask", visible=False)
+            gr.ClearButton(components=[input_img, output_img, mask_img])
+
+    chk_include_matting.change(change_include_matting, inputs=[chk_include_matting],
+                               outputs=[slider_matting, matting_state,
+                                        slr_fg_threshold, slr_bg_threshold, slr_erode_size])
+
+    slr_fg_threshold.change(change_foreground_threshold, inputs=[slr_fg_threshold, matting_state],
+                            outputs=[matting_state])
+
+    slr_bg_threshold.change(change_background_threshold, inputs=[slr_bg_threshold, matting_state],
+                            outputs=[matting_state])
+
+    slr_erode_size.change(change_erode_size, inputs=[slr_erode_size, matting_state],
+                          outputs=[matting_state])
+
+    chk_show_mask.change(change_show_mask, inputs=[chk_show_mask], outputs=[mask_img])
+
+    chk_change_color.change(change_background_mode, inputs=[chk_change_color],
+                            outputs=[pkr_color, chk_dominant])
+
+    pkr_color.change(change_picker_color, inputs=[pkr_color, chk_dominant], outputs=[color_state])
 
-    chk_alm.change(alpha_matting, inputs=[chk_alm], outputs=[sld_aft, sld_amb, sld_aes])
+    chk_dominant.change(set_dominant_color, inputs=[chk_dominant], outputs=[color_state, pkr_color])
 
-    run_btn.click(predict, [input_img, drp_models,
-                            chk_alm, chk_msk, chk_psm,
-                            sld_aft, sld_amb, sld_aes, chk_col, color, trans], [output_img])
+    run_btn.click(predict, inputs=[input_img, drp_models, chk_smoot_mask, matting_state, color_state],
+                  outputs=[output_img, mask_img])
 
     with gr.Row():
         gr.HTML(footer)

cutter.py

@@ -0,0 +1,98 @@
+import PIL
+import numpy as np
+from PIL import Image, ImageColor, ImageDraw
+from PIL.Image import Image as PILImage
+from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
+from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
+from pymatting.util.util import stack_images
+from rembg.bg import post_process, naive_cutout, apply_background_color
+from scipy.ndimage import binary_erosion
+
+
+def alpha_matting_cutout(img: PILImage, trimap: np.ndarray) -> PILImage:
+    if img.mode == "RGBA" or img.mode == "CMYK":
+        img = img.convert("RGB")
+
+    img = np.asarray(img)
+
+    img_normalized = img / 255.0
+    trimap_normalized = trimap / 255.0
+
+    alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
+    foreground = estimate_foreground_ml(img_normalized, alpha)
+    cutout = stack_images(foreground, alpha)
+
+    cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)
+    return Image.fromarray(cutout)
+
+
+def generate_trimap(
+        mask: PILImage,
+        foreground_threshold: int,
+        background_threshold: int,
+        erode_structure_size: int,
+) -> np.ndarray:
+    mask = np.asarray(mask)
+
+    is_foreground = mask > foreground_threshold
+    is_background = mask < background_threshold
+
+    structure = None
+    if erode_structure_size > 0:
+        structure = np.ones(
+            (erode_structure_size, erode_structure_size), dtype=np.uint8
+        )
+
+    is_foreground = binary_erosion(is_foreground, structure=structure)
+    is_background = binary_erosion(is_background, structure=structure, border_value=1)
+
+    trimap = np.full(mask.shape, dtype=np.uint8, fill_value=128)
+    trimap[is_foreground] = 255
+    trimap[is_background] = 0
+
+    return trimap
+
+
+def get_background_dominant_color(img: PILImage, mask: PILImage) -> tuple:
+    negative_img = img.copy()
+    negative_mask = PIL.ImageOps.invert(mask)
+    negative_img.putalpha(negative_mask)
+    negative_img = negative_img.resize((1, 1))
+    r, g, b, a = negative_img.getpixel((0, 0))
+    return r, g, b, 255
+
+
+def remove(session, img: PILImage, smoot: bool, matting: tuple, color) -> (PILImage, PILImage):
+    mask = session.predict(img)[0]
+
+    if smoot:
+        mask = PIL.Image.fromarray(post_process(np.array(mask)))
+
+    fg_t, bg_t, erode = matting
+
+    if fg_t > 0 or bg_t > 0 or erode > 0:
+        mask = generate_trimap(mask, *matting)
+        try:
+            cutout = alpha_matting_cutout(img, mask)
+            mask = PIL.Image.fromarray(mask)
+        except ValueError as err:
+            raise err
+    else:
+        cutout = naive_cutout(img, mask)
+
+    if color is True:
+        color = get_background_dominant_color(img, mask)
+        cutout = apply_background_color(cutout, color)
+    elif isinstance(color, str):
+        r, g, b = ImageColor.getcolor(color, "RGB")
+        cutout = apply_background_color(cutout, (r, g, b, 255))
+
+    return cutout, mask
+
+
+def make_label(text, width=600, height=200, color="black") -> PILImage:
+    image = Image.new("RGB", (width, height), color)
+    draw = ImageDraw.Draw(image)
+    text_width, text_height = draw.textsize(text)
+    draw.text(((width-text_width)/2, height/2), text)
+    return image

requirements.txt

@@ -1,3 +1,4 @@
 rembg~=2.0.47
 pillow~=9.5.0
-opencv-python-headless
+pymatting
+opencv-python-headless