Update app.py

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-from gradio_imageslider import ImageSlider
 from loadimg import load_img
 import spaces
 from transformers import AutoModelForImageSegmentation
@@ -8,19 +7,19 @@ from torchvision import transforms
 import uuid
 import os
 
-# Automatically select device based on availability
+# Automatically select device (GPU or CPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Optional: Print which device is being used
 print(f"Using device: {device}")
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
+# Load the segmentation model
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
 birefnet.to(device)
 
+# Image preprocessing
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -33,7 +32,6 @@ transform_image = transforms.Compose(
 def process(image):
     image_size = image.size
     input_images = transform_image(image).unsqueeze(0).to(device)
-    # Prediction
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
@@ -42,47 +40,45 @@ def process(image):
     image.putalpha(mask)
     return image
 
+# Main function: returns both a preview and a downloadable file
 def fn(image):
-    im = load_img(image, output_type="pil")
-    im = im.convert("RGB")
+    im = load_img(image, output_type="pil").convert("RGB")
     processed_image = process(im)
 
-    # Save to a temp file
     filename = f"/tmp/processed_{uuid.uuid4().hex}.png"
     processed_image.save(filename)
 
-    return processed_image, filename  # Return both preview and downloadable file
+    return processed_image, filename  # Preview + downloadable image
 
+# File-only processing tab
 def process_file(f):
     name_path = f.rsplit(".", 1)[0] + ".png"
-    im = load_img(f, output_type="pil")
-    im = im.convert("RGB")
+    im = load_img(f, output_type="pil").convert("RGB")
     transparent = process(im)
     transparent.save(name_path)
     return name_path
 
 # Components
-slider1 = gr.Image(label="Preview")  # Changed from ImageSlider to Image
-output_file1 = gr.File(label="Download Processed Image")
-slider2 = ImageSlider(label="Processed Image from URL", type="pil")
+preview_image = gr.Image(label="Preview")
+download_file = gr.File(label="Download Processed Image")
 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("dog.jpeg", output_type="pil")
+# Example image (from URL instead of invalid local path)
 url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
+chameleon = load_img(url_example, output_type="pil")
 
-# Tabs
-tab1 = gr.Interface(fn, inputs=image_upload, outputs=[slider1, output_file1], 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=["dog.jpeg"], api_name="png")
+# Interface Tabs
+tab1 = gr.Interface(fn, inputs=image_upload, outputs=[preview_image, download_file], examples=[chameleon], api_name="image")
+tab2 = gr.Interface(fn, inputs=url_input, outputs=preview_image, examples=[url_example], api_name="text")
+tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=[], api_name="png")
 
+# Combine tabs into one app
 demo = gr.TabbedInterface(
     [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
 )
 
-
 if __name__ == "__main__":
     demo.launch(show_error=True)