Update app.py

app.py

@@ -3,7 +3,7 @@ import cv2
 import numpy as np
 import torch
 import gradio as gr
-import spaces
+import spaces  # Added import for spaces
 
 from PIL import Image, ImageOps
 from transformers import AutoModelForImageSegmentation
@@ -56,18 +56,19 @@ class ImagePreprocessor():
         image = self.transform_image(image)
         return image
 
-birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet-matting', trust_remote_code=True)
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    'zhengpeng7/BiRefNet-matting', trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
-@spaces.GPU
+@spaces.GPU  # Added the @spaces.GPU decorator
 def remove_background(image):
     if image is None:
         raise gr.Error("Please upload an image.")
 
     image_ori = Image.fromarray(image).convert('RGB')
     original_size = image_ori.size
-    
+
     # Preprocess the image
     image_preprocessor = ImagePreprocessor(resolution=(1024, 1024))
     image_proc = image_preprocessor.proc(image_ori)
@@ -81,10 +82,9 @@ def remove_background(image):
     # Process Results
     pred_pil = transforms.ToPILImage()(pred)
     pred_pil = pred_pil.resize(original_size, Image.BICUBIC)  # Resize mask to original size
-    
-    # Create reverse mask
-    reverse_mask = Image.new('L', original_size)
-    reverse_mask.paste(ImageOps.invert(pred_pil))
+
+    # Create reverse mask (background mask)
+    reverse_mask = ImageOps.invert(pred_pil)
 
     # Create foreground image (object with transparent background)
     foreground = image_ori.copy()
@@ -96,29 +96,17 @@ def remove_background(image):
 
     torch.cuda.empty_cache()
 
-    # Save all images
-    mask_path = "mask.png"
-    pred_pil.save(mask_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, reverse_mask_path, foreground_path, background_path
+    # Return images in the specified order
+    return foreground, background, pred_pil, reverse_mask
 
 iface = gr.Interface(
     fn=remove_background,
     inputs=gr.Image(type="numpy"),
     outputs=[
-        gr.Image(type="filepath", label="Mask"),
-        gr.Image(type="filepath", label="Reverse Mask"),
-        gr.Image(type="filepath", label="Foreground"),
-        gr.Image(type="filepath", label="Background")
+        gr.Image(type="pil", label="Foreground"),
+        gr.Image(type="pil", label="Background"),
+        gr.Image(type="pil", label="Foreground Mask"),
+        gr.Image(type="pil", label="Background Mask")
     ],
     allow_flagging="never"
 )