Update app.py

app.py

@@ -67,10 +67,13 @@ def extract_image(img, pos_prompts, neg_prompts, threshold, alpha_value=0.5):
     output_image[:, :, :3] = img_np[:, :, :3]
     output_image[:, :, 3] = (final_mask * alpha_value * 255).astype(np.uint8)
 
-    # convert the output_image back to a PIL.Image object
+    # 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
 
-    return output_image, final_mask.astype(np.uint8), bmask