Update app.py

app.py

@@ -1,6 +1,6 @@
 from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
 import gradio as gr
-from PIL import Image, ImageFilter
+from PIL import Image
 import torch
 import matplotlib.pyplot as plt
 import torch
@@ -44,28 +44,33 @@ def get_masks(prompts, img, threhsold):
     return masks
 
 
-def extract_image(img, pos_prompts, neg_prompts, threshold, blur_radius=5):
+def extract_image(img, pos_prompts, neg_prompts, threshold, alpha_value=0.5):
     positive_masks = get_masks(pos_prompts, img, threshold)
     negative_masks = get_masks(neg_prompts, img, threshold)
 
-    # combine masks into one masks, logic OR
+    # combine masks into one mask, logic OR
     pos_mask = np.any(np.stack(positive_masks), axis=0)
     neg_mask = np.any(np.stack(negative_masks), axis=0)
     final_mask = pos_mask & ~neg_mask
 
-    # apply Gaussian blur for feathering
-    final_mask_img = Image.fromarray((final_mask * 255).astype(np.uint8), "L")
-    final_mask_img = final_mask_img.filter(ImageFilter.GaussianBlur(radius=blur_radius))
-    final_mask = np.array(final_mask_img) / 255
-    final_mask = final_mask > threshold
+    # threshold the mask
+    bmask = final_mask > threshold
+    # zero out values below the threshold
+    final_mask[final_mask < threshold] = 0
 
-    # extract the final image
-    final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255, "L")
-    inverse_mask = np.invert(final_mask)
-    output_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
-    output_image.paste(img, mask=final_mask)
+    # convert PIL image to RGBA numpy array
+    img_np = np.array(img.convert("RGBA"))
+    # create an empty RGBA image with the same size
+    output_image = np.zeros_like(img_np)
 
-    return output_image, final_mask, inverse_mask
+    # apply the final_mask as alpha channel on the output image
+    output_image[:, :, :3] = img_np[:, :, :3]
+    output_image[:, :, 3] = (final_mask * 255 * alpha_value).astype(np.uint8)
+
+    # convert the output_image back to a PIL.Image object
+    output_image = Image.fromarray(output_image, "RGBA")
+
+    return output_image, final_mask, bmask