Update app.py

app.py

@@ -1,59 +1,68 @@
 import gradio as gr
-from PIL import Image
 import torch
 import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+from u2net import U2NET
 import os
-from torchvision import transforms
-from u2net import U2NET # Load model class
+import urllib.request
 
-# Load the model (download once, reuse)
+# 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):
-    import requests
-    url = "https://huggingface.co/akhaliq/U-2-Net/resolve/main/u2net.pth"
-    with open(model_path, "wb") as f:
-        f.write(requests.get(url).content)
-
-# Load model to CPU
-net = U2NET(3,1)
-net.load_state_dict(torch.load(model_path, map_location='cpu'))
+    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
-transform = transforms.Compose([
-    transforms.Resize((320,320)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225])
-])
-
-# Post-process output mask
-def normalize_prediction(pred):
-    ma = torch.max(pred)
-    mi = torch.min(pred)
-    return (pred - mi) / (ma - mi)
-
-# Main background remover function
-def remove_bg(input_image):
-    image = input_image.convert("RGB")
-    orig_size = image.size
-    img_tensor = transform(image).unsqueeze(0)
+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(img_tensor)
-        mask = normalize_prediction(d1[0][0])
-        mask = mask.squeeze().cpu().numpy()
-        mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(orig_size)
-
-    # Add alpha channel using mask
-    image.putalpha(mask)
-    return image
-
-# Gradio app
-gr.Interface(
-    fn=remove_bg,
+        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="Transparent PNG"),
-    title="🪄 Remove.bg Clone - AI Background Remover",
-    description="Upload any image to remove the background using U-2-Net (Hugging Face version)."
-).launch()
+    outputs=gr.Image(type="pil", label="Image without Background"),
+    title="🧠 AI Background Remover (U²-Net)",
+    description="Removes background from images using U²-Net."
+)
+
+demo.launch()