Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import AutoModelForImageSegmentation
 from PIL import Image
 from torchvision import transforms
 import gradio as gr
+import os
 
 # Load the model from Hugging Face
 birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet_lite', trust_remote_code=True)
@@ -11,7 +12,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 birefnet.to(device)
 birefnet.eval()
 
-# Define the transform to preprocess the input image
+# Define the transform to preprocess the input images
 image_size = (1024, 1024)
 transform_image = transforms.Compose([
     transforms.Resize(image_size),
@@ -19,10 +20,11 @@ transform_image = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-def extract_objects(images):
-    start_time = time.time()
+def extract_objects(files):
+    # Open all images from the uploaded files
+    images = [Image.open(file.name).convert("RGB") for file in files]
 
-    # Transform all images into a batch
+    start_time = time.time()
     inputs = []
     original_sizes = []
     for img in images:
@@ -70,10 +72,10 @@ def extract_objects(images):
 
 iface = gr.Interface(
     fn=extract_objects,
-    inputs=gr.Image(type="pil", label="Upload Images", multiple=True),
+    inputs=gr.Files(label="Upload Multiple Images", type="file", file_count="multiple"),
     outputs=[gr.Gallery(label="Processed Images"), gr.Textbox(label="Timing Info")],
     title="BiRefNet Bulk Background Removal",
-    description="Upload multiple images and process them in bulk. The request is handled at once, not sequentially. Timing information is also provided."
+    description="Upload multiple images and process them in one request. Timing information for the full request and per-image processing is provided."
 )
 
 if __name__ == "__main__":