Update app.py

app.py

@@ -40,11 +40,16 @@ def get_masks(prompts, img, threhsold):
     for prompt in prompts:
         mask = process_image(img, prompt)
         mask = mask > threhsold
+        ax = plt.subplots()
+        ax.imshow(image)
+        ax.imshow(mask, alpha=0.5, cmap="jet")
+        ax.axis("off")
+        plt.tight_layout()
         masks.append(mask)
     return masks
 
 
-def extract_image(img, pos_prompts, neg_prompts, threshold, alpha_value=0.5):
+def extract_image(img, pos_prompts, neg_prompts, threshold):
     positive_masks = get_masks(pos_prompts, img, threshold)
     negative_masks = get_masks(neg_prompts, img, threshold)
 
@@ -53,27 +58,13 @@ 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)
-
-    # apply alpha value to final_mask
-    alpha_mask = Image.fromarray((final_mask * 255 * alpha_value).astype(np.uint8), "L")
-    inverse_alpha_mask = Image.fromarray(((1 - final_mask) * 255 * alpha_value).astype(np.uint8), "L")
-
     # 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=alpha_mask)
-
-    # extract the inverse_mask
-    output_inverse_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
-    output_inverse_image.paste(img, mask=inverse_alpha_mask)
+    output_image.paste(img, mask=final_mask)
 
-    return output_image, alpha_mask, inverse_alpha_mask
+    return output_image, final_mask, inverse_mask
 
 
 title = "Interactive demo: zero-shot image segmentation with CLIPSeg"