initial upload

app.py

@@ -0,0 +1,66 @@
+import torch
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+from PIL import Image
+import requests
+from io import BytesIO
+import gradio as gr
+# Set up CUDA if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_float32_matmul_precision("high")
+
+# Load the model
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+birefnet.to(device)
+
+# Define image transformations
+transform_image = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+
+def load_img(image_path_or_url):
+    if image_path_or_url.startswith('http'):
+        response = requests.get(image_path_or_url)
+        img = Image.open(BytesIO(response.content))
+    else:
+        img = Image.open(image_path_or_url)
+    return img.convert("RGB")
+
+def process(image):
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to(device)
+
+    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)
+
+    # Create a new image with transparency
+    transparent_image = Image.new("RGBA", image.size)
+    transparent_image.paste(image, (0, 0))
+    transparent_image.putalpha(mask)  # Apply mask to the new image
+
+    return transparent_image  # Return the new transparent image
+
+def remove_background_gradio(image):
+    processed_img = process(image)
+    return processed_img
+
+
+# Create the Gradio interface with drag-and-drop and paste functionality
+demo = gr.Interface(
+    fn=remove_background_gradio,
+    inputs=gr.Image(type="pil"),  # Remove 'source' argument
+    outputs=gr.Image(type="pil"),
+    title="RemoveBG",
+    description="Upload an image to remove its background (drag-and-drop or upload)."
+)
+
+demo.launch(share=True)  # Launch the interface and get a shareable link

requirements.txt

@@ -0,0 +1,8 @@
+transformers
+torch
+torchvision
+Pillow
+requests
+kornia
+timm
+gradio