Add alpha matting

app.py

@@ -6,6 +6,9 @@ from PIL import Image, ImageDraw, ImageFilter
 import numpy as np
 import spaces
 from huggingface_hub import hf_hub_download
+from trimap_module import trimap
+from pymatting import *
+
 
 pipe = FluxFillPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Fill-dev",
@@ -140,12 +143,27 @@ def inpaint(image, width, height, overlap_percentage, num_inference_steps, resiz
         guidance_scale=30,
     ).images[0]
    
-    result = result.convert("RGBA")
+    result = result.convert("RGB")
     # mask_blur = mask.filter(ImageFilter.GaussianBlur(10))
-    # cnet_image = Image.composite(result, cnet_image, mask) 
-    cnet_image.paste(result, (0, 0), mask)
+    # cnet_image = Image.composite(result, cnet_image, mask)
+    size    = 10; # how many pixel extension do you want to dilate
+    number  = 1;  # numbering purpose 
+    trimap_mask = trimap(image, name, size, number, erosion=False)
+
+    # estimate alpha from image and trimap
+    alpha = estimate_alpha_cf(cnet_image, trimap_mask)
+    
+    # make gray background
+    background = result
+    
+    # estimate foreground from image and alpha
+    foreground = estimate_foreground_ml(cnet_image, alpha)
+    
+    # blend foreground with background and alpha, less color bleeding
+    new_image = blend(foreground, background, alpha)
+    # cnet_image.paste(result, (0, 0), mask)
 
-    return cnet_image, background
+    return new_image, background
 
 def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)