Update app.py

app.py

@@ -1,3 +1,4 @@
+import time
 import torch
 from transformers import AutoModelForImageSegmentation
 from PIL import Image
@@ -18,25 +19,61 @@ transform_image = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-def extract_object(image):
-    input_images = transform_image(image).unsqueeze(0).to(device)
+def extract_objects(images):
+    start_time = time.time()
+
+    # Transform all images into a batch
+    inputs = []
+    original_sizes = []
+    for img in images:
+        original_sizes.append(img.size)
+        inputs.append(transform_image(img))
+    input_tensor = torch.stack(inputs).to(device)
+
+    # Inference
+    inf_start = time.time()
     with torch.no_grad():
-        preds = birefnet(input_images)[-1].sigmoid().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image.size)
+        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+    inf_end = time.time()
+
+    # Post-process results
+    results = []
+    image_times = []
+    for i, img in enumerate(images):
+        t_start = time.time()
+        pred = preds[i].squeeze()
+        pred_pil = transforms.ToPILImage()(pred)
+        mask = pred_pil.resize(original_sizes[i])
+        
+        # Create a transparent background image
+        result = Image.new("RGBA", original_sizes[i], (0, 0, 0, 0))
+        result.paste(img, mask=mask)
+        results.append(result)
+        t_end = time.time()
+        image_times.append(t_end - t_start)
+
+    end_time = time.time()
+    total_time = end_time - start_time
+    inference_time = inf_end - inf_start
+    prep_post_time = total_time - inference_time
+
+    # Create a summary of timings
+    summary = (
+        f"Total request time: {total_time:.2f} s\n"
+        f"Inference time (batch): {inference_time:.2f} s\n"
+        f"Pre/Post-processing time: {prep_post_time:.2f} s\n"
+        "Per-image post-processing times:\n" + 
+        "\n".join([f"  Image {i+1}: {t:.2f} s" for i, t in enumerate(image_times)])
+    )
 
-    # Create a new transparent image
-    result = Image.new("RGBA", image.size, (0, 0, 0, 0))
-    result.paste(image, mask=mask)
-    return result
+    return results, summary
 
 iface = gr.Interface(
-    fn=extract_object,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="pil", label="Object with Transparent Background"),
-    title="BiRefNet Background Removal",
-    description="Upload an image and get the foreground object extracted onto a transparent background."
+    fn=extract_objects,
+    inputs=gr.Image(type="pil", label="Upload Images", multiple=True),
+    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."
 )
 
 if __name__ == "__main__":