Update app.py

app.py

@@ -0,0 +1,61 @@
+# app.py
+import gradio as gr
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+import numpy as np
+import sys
+
+# Install and import (for Hugging Face environment)
+try:
+    import pytorch_hub_examples
+except ModuleNotFoundError:  # Handle if not installed (Hugging Face)
+    import subprocess
+    subprocess.run(["git", "clone", "https://github.com/facebookresearch/pytorch_hub_examples.git"])
+    subprocess.run(["cd", "pytorch_hub_examples", "&&", "pip", "install", "-e", "."])
+    import pytorch_hub_examples
+
+# Load the U-2-Net model
+model = pytorch_hub_examples.u2net(pretrained=True)
+model.eval()
+
+# Define the transform
+transform = transforms.Compose([
+    transforms.Resize((320, 320)),  # Resize for U-2-Net
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ImageNet stats
+])
+
+
+def remove_background(image):
+    try:
+        img = transform(image).unsqueeze(0)
+
+        with torch.no_grad():
+            out = model(img)
+            mask = torch.sigmoid(out[0]) # Get the mask
+
+        mask = mask.squeeze().cpu().numpy()
+        mask = (mask * 255).astype(np.uint8)
+        mask = Image.fromarray(mask).convert("L")
+
+        image = image.convert("RGBA")
+        new_image = Image.new("RGBA", image.size, (255, 255, 255, 0))
+
+        for x in range(image.width):
+            for y in range(image.height):
+                if mask.getpixel((x, y)) > 0:
+                    new_image.putpixel((x, y), image.getpixel((x, y)))
+
+        return new_image
+    except Exception as e:
+        return f"Error: {e}"
+
+
+with gr.Blocks() as demo:
+    image_input = gr.Image(type="pil")
+    image_output = gr.Image()
+    image_button = gr.Button("Remove Background")
+    image_button.click(remove_background, inputs=image_input, outputs=image_output)
+
+demo.launch()