Update app.py

app.py

@@ -44,36 +44,35 @@ def get_masks(prompts, img, threhsold):
     return masks
 
 
-def extract_image(img, pos_prompts, neg_prompts, threshold, alpha_value=0.5):
+def extract_image(img, pos_prompts, neg_prompts, threshold, alpha_value=0.5, blur_radius=5):
     positive_masks = get_masks(pos_prompts, img, threshold)
     negative_masks = get_masks(neg_prompts, img, threshold)
 
-    # combine masks into one mask, logic OR
+    # combine masks into one masks, 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
 
-    # threshold the mask
-    bmask = final_mask > threshold
-    # zero out values below the threshold
-    final_mask[final_mask < threshold] = 0
-
-    # 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, dtype=np.uint8)
-
-    # apply the final_mask as alpha channel on the output image
-    output_image[:, :, :3] = img_np[:, :, :3]
-    output_image[:, :, 3] = (final_mask * alpha_value * 255).astype(np.uint8)
-
-    # convert the output_image, final_mask, and bmask back to PIL.Image objects
-    output_image = Image.fromarray(output_image, "RGBA")
-    final_mask = Image.fromarray((final_mask * 255).astype(np.uint8), "L")
-    bmask = Image.fromarray((bmask * 255).astype(np.uint8), "L")
-
-    return output_image, final_mask, bmask
-
+    # 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
+
+    # blend the original image and the mask using the alpha value
+    fig, ax = plt.subplots()
+    ax.imshow(img)
+    ax.imshow(final_mask, alpha=alpha_value, cmap="jet")
+    ax.axis("off")
+    plt.tight_layout()
+
+    # 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)
+
+    return output_image, final_mask, inverse_mask