Update app.py

app.py

@@ -1,68 +1,28 @@
 import gradio as gr
-import torch
-import numpy as np
 from PIL import Image
-import torchvision.transforms as transforms
-from u2net import U2NET
-import os
-import urllib.request
+from rembg import remove
+import io
 
-# Download model if not present
-model_url = "https://huggingface.co/flashingtt/U-2-Net/resolve/main/u2net.pth"
-model_path = "u2net.pth"
-
-if not os.path.exists(model_path):
-    print("Downloading model...")
-    urllib.request.urlretrieve(model_url, model_path)
-
-# Load model
-print("Loading model...")
-net = U2NET(3, 1)
-net.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=False))
-net.eval()
-
-# Preprocessing
-def preprocess(img):
-    transform = transforms.Compose([
-        transforms.Resize((320, 320)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                             std=[0.229, 0.224, 0.225])
-    ])
-    return transform(img).unsqueeze(0)
-
-# Postprocess output mask
-def postprocess_mask(d):
-    pred = d[0][0]
-    pred = (pred - pred.min()) / (pred.max() - pred.min())
-    pred = pred.detach().cpu().numpy()
-    mask = (pred > 0.5).astype(np.uint8) * 255
-    return Image.fromarray(mask)
-
-# Main function
 def remove_background(input_image):
-    image = input_image.convert("RGB")
-    input_tensor = preprocess(image)
-    with torch.no_grad():
-        d1, *_ = net(input_tensor)
-    mask = postprocess_mask(d1)
-
-    image = image.resize(mask.size)
-    image_np = np.array(image)
-    mask_np = np.array(mask) / 255
-    mask_np = np.expand_dims(mask_np, axis=2)
-
-    result = image_np * mask_np + (1 - mask_np) * 255
-    result = Image.fromarray(result.astype(np.uint8))
-    return result
-
-# Gradio UI
-demo = gr.Interface(
-    fn=remove_background,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="pil", label="Image without Background"),
-    title="🧠 AI Background Remover (U²-Net)",
-    description="Removes background from images using U²-Net."
-)
+    # Ensure image is RGBA for transparency
+    input_image = input_image.convert("RGBA")
+    
+    # Remove background
+    result = remove(input_image)
+    
+    return Image.open(io.BytesIO(result))
+
+with gr.Blocks(title="AI Background Remover (rembg)") as demo:
+    gr.Markdown("🧠 **AI Background Remover** - Powered by `rembg` for transparent results")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Upload Image", type="pil")
+            submit = gr.Button("Remove Background")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Transparent PNG Output")
+
+    submit.click(fn=remove_background, inputs=input_image, outputs=output_image)
 
 demo.launch()