Update app.py

app.py

@@ -5,7 +5,7 @@ import torch
 import gradio as gr
 import spaces
 
-from PIL import Image
+from PIL import Image, ImageOps
 from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 
@@ -78,23 +78,38 @@ def remove_background(image):
         preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
 
-    # Show Results
+    # Process Results
     pred_pil = transforms.ToPILImage()(pred)
     pred_pil = pred_pil.resize(original_size, Image.BICUBIC)  # Resize mask to original size
-    image_masked = refine_foreground(image_ori, pred_pil)
-    image_masked.putalpha(pred_pil)
+    
+    # Create reverse mask
+    reverse_mask = Image.new('L', original_size)
+    reverse_mask.paste(ImageOps.invert(pred_pil))
+
+    # Create foreground image (object with transparent background)
+    foreground = image_ori.copy()
+    foreground.putalpha(pred_pil)
+
+    # Create background image
+    background = image_ori.copy()
+    background.putalpha(reverse_mask)
 
     torch.cuda.empty_cache()
 
-    # Save mask as PNG
+    # Save all images
     mask_path = "mask.png"
     pred_pil.save(mask_path)
 
-    # Save output as PNG
-    output_path = "output.png"
-    image_masked.save(output_path)
+    reverse_mask_path = "reverse_mask.png"
+    reverse_mask.save(reverse_mask_path)
+
+    foreground_path = "foreground.png"
+    foreground.save(foreground_path)
+
+    background_path = "background.png"
+    background.save(background_path)
 
-    return mask_path, output_path
+    return mask_path, reverse_mask_path, foreground_path, background_path
 
 css = """
 body {
@@ -139,7 +154,9 @@ iface = gr.Interface(
     inputs=gr.Image(type="numpy"),
     outputs=[
         gr.Image(type="filepath", label="Mask"),
-        gr.Image(type="filepath", label="Output")
+        gr.Image(type="filepath", label="Reverse Mask"),
+        gr.Image(type="filepath", label="Foreground"),
+        gr.Image(type="filepath", label="Background")
     ],
     allow_flagging="never",
     css=css