Update app.py

app.py

@@ -53,20 +53,23 @@ def extract_image(img, pos_prompts, neg_prompts, threshold, alpha_value=0.5):
     neg_mask = np.any(np.stack(negative_masks), axis=0)
     final_mask = pos_mask & ~neg_mask
 
+    # create the final mask image
+    final_mask_img = Image.fromarray((final_mask * 255).astype(np.uint8), "L")
+
+    # create an RGBA version of the original image and the final mask image
+    img_rgba = img.convert("RGBA")
+    mask_rgba = Image.new("RGBA", img.size, (0, 0, 0, 0))
+    mask_rgba.paste(final_mask_img, (0, 0), final_mask_img)
+
     # 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()
+    output_image = Image.blend(img_rgba, mask_rgba, alpha_value)
 
     # 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, mask, inverse_mask
+    return output_image, final_mask, inverse_mask
+
 
 
 title = "Interactive demo: zero-shot image segmentation with CLIPSeg"