updated download function

app.py

@@ -1,17 +1,27 @@
 import gradio as gr
-from loadimg import load_img
+from PIL import Image
 import spaces
 from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
+import requests
+from io import BytesIO
+import os
 
-torch.set_float32_matmul_precision(["high", "highest"][0])
+# --- Model and Processor Setup ---
+# Use a higher precision for matrix multiplication for better performance
+torch.set_float32_matmul_precision("high")
 
+# Load the BiRefNet model for image segmentation
+# trust_remote_code=True is required for this model
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-birefnet.to("cpu")
+# Move the model to the available device (GPU if available, otherwise CPU)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+birefnet.to(device)
 
+# Define the image transformation pipeline
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -20,52 +30,144 @@ transform_image = transforms.Compose(
     ]
 )
 
-def fn(image):
-    im = load_img(image, output_type="pil")
-    im = im.convert("RGB")
-    origin = im.copy()
-    processed_image = process(im)
-    return (processed_image, origin)
+# --- Helper Function to Load Images ---
+def load_image(image_source, output_type="pil"):
+    """
+    Loads an image from a file path, URL, or numpy array.
+    """
+    if image_source is None:
+        return None
+        
+    if isinstance(image_source, str):
+        if image_source.startswith("http"):
+            try:
+                response = requests.get(image_source)
+                response.raise_for_status()
+                image = Image.open(BytesIO(response.content))
+            except requests.exceptions.RequestException as e:
+                raise gr.Error(f"Could not fetch image from URL: {e}")
+        else:
+            image = Image.open(image_source)
+    elif hasattr(image_source, 'shape'): # Check if it's a numpy-like array
+        image = Image.fromarray(image_source)
+    else:
+        image = image_source # Assume it's already a PIL image
 
-#@spaces.GPU
-def process(image):
+    if output_type == "pil":
+        return image.convert("RGB")
+    return image
+
+# --- Core Processing Function ---
+# Use @spaces.GPU decorator if you plan to run this on a GPU-enabled Hugging Face Space
+# @spaces.GPU
+def process_image_to_transparent(image: Image.Image) -> Image.Image:
+    """
+    Takes a PIL image, removes the background, and returns a PIL image with an alpha channel.
+    """
+    if image is None:
+        return None
     image_size = image.size
-    input_images = transform_image(image).unsqueeze(0).to("cpu")
-    # Prediction
+    # Unsqueeze adds a batch dimension, which the model expects
+    input_tensor = transform_image(image).unsqueeze(0).to(device)
+
+    # Prediction without tracking gradients for efficiency
     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)
-    image.putalpha(mask)
+        # The model returns multiple outputs; the last one is the primary segmentation map
+        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+
+    # Process the prediction tensor to create a mask
+    pred_tensor = preds[0].squeeze()
+    mask_pil = transforms.ToPILImage()(pred_tensor)
+    mask_resized = mask_pil.resize(image_size)
+
+    # Apply the mask as an alpha channel to the original image
+    image.putalpha(mask_resized)
     return image
 
-def process_file(f):
-    name_path = f.rsplit(".", 1)[0] + ".png"
-    im = load_img(f, output_type="pil")
-    im = im.convert("RGB")
-    transparent = process(im)
-    transparent.save(name_path)
-    return name_path
-
-slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png")
-slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png")
-image_upload = gr.Image(label="Upload an image")
-image_file_upload = gr.Image(label="Upload an image", type="filepath")
-url_input = gr.Textbox(label="Paste an image URL")
-output_file = gr.File(label="Output PNG File")
-
-# Example images
-chameleon = load_img("butterfly.jpeg", output_type="pil") 
-url_example = "https://i.ibb.co/67B6Knk9/students-1807505-1280.jpg"
-
-tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
-tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
-tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpeg"], api_name="png")
+# --- Gradio Interface Functions ---
+
+def fn(image_source):
+    """
+    Handles image uploads and URLs, returning the processed image.
+    """
+    if image_source is None:
+        return None
+    
+    pil_image = load_image(image_source, output_type="pil")
+    processed_image = process_image_to_transparent(pil_image)
+    return processed_image
+
+def process_file(image_filepath):
+    """
+    Handles a single file upload and returns a downloadable processed file.
+    """
+    if image_filepath is None:
+        return None
+
+    # Define the output path for the new PNG file
+    base_name = os.path.basename(image_filepath.name) # Use .name for Gradio file objects
+    name, _ = os.path.splitext(base_name)
+    output_path = f"{name}_transparent.png"
+    
+    # Load the image from the provided file path
+    pil_image = load_image(image_filepath.name, output_type="pil")
+    
+    # Process the image
+    transparent_image = process_image_to_transparent(pil_image)
+    
+    # Save the processed image to the new path
+    transparent_image.save(output_path)
+    
+    # Return the path to the newly created file for download
+    return output_path
+
+# --- Gradio UI Definition ---
+
+# Define example images for the interface
+example_image_path = "butterfly.jpeg"
+# You should have a 'butterfly.jpeg' in the same directory or provide a full path
+# For demonstration, let's create a dummy example image if it doesn't exist.
+if not os.path.exists(example_image_path):
+    print(f"'{example_image_path}' not found. Creating a dummy image for example.")
+    try:
+        dummy_img = Image.new('RGB', (200, 200), color = 'red')
+        dummy_img.save(example_image_path)
+    except Exception as e:
+        print(f"Could not create dummy image: {e}")
+
+example_url = "https://i.ibb.co/67B6Knk9/students-1807505-1280.jpg"
+
+# Define the individual interfaces for each tab
+tab1 = gr.Interface(
+    fn, 
+    inputs=gr.Image(label="Upload an Image", type="pil"), 
+    outputs=gr.Image(label="Processed Image", format="png"), 
+    examples=[[example_image_path]], 
+    api_name="image"
+)
+
+tab2 = gr.Interface(
+    fn, 
+    inputs=gr.Textbox(label="Paste an Image URL"), 
+    outputs=gr.Image(label="Processed Image", format="png"), 
+    examples=[[example_url]], 
+    api_name="text"
+)
+
+tab3 = gr.Interface(
+    process_file, 
+    inputs=gr.File(label="Upload an Image File"), 
+    outputs=gr.File(label="Download Processed PNG"), 
+    examples=[[example_image_path]], 
+    api_name="png"
+)
 
+# Combine the interfaces into a tabbed layout
 demo = gr.TabbedInterface(
-    [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
+    [tab1, tab2, tab3], 
+    ["Image Upload", "URL Input", "File Output"], 
+    title="Background Removal Tool"
 )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(show_error=True)