Upload 4 files

README.md

@@ -1,13 +1,104 @@
----
-title: Background Removal Tool
-emoji: 👀
-colorFrom: blue
-colorTo: blue
-sdk: gradio
-sdk_version: 5.12.0
-app_file: app.py
-pinned: false
-short_description: A tool to remove image backgrounds with precision
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Background Removal Tool
+emoji: 👀
+colorFrom: blue
+colorTo: blue
+sdk: gradio
+sdk_version: 5.12.0
+app_file: app.py
+pinned: false
+short_description: A tool to remove image backgrounds with precision
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Background Removal Tool
+
+This is a deep learning-powered **Background Removal Tool** that uses image segmentation models to remove backgrounds from images and add transparency (alpha channel). It features a user-friendly interface built with Gradio to interact with the tool via image uploads, URLs, or file outputs.
+
+---
+
+## Features
+
+1. **Two Segmentation Models**:
+   - `BiRefNet`: Efficient and robust segmentation model.
+   - `RMBG-2.0`: Advanced model for refined background removal.
+2. **Multiple Input Methods**:
+   - Upload images directly from your system.
+   - Provide an image URL for processing.
+   - Upload and save the processed image as a PNG file with transparency.
+3. **Customizable**: Switch between models for different use cases.
+4. **Fast and GPU-Powered**: Leverages CUDA for faster processing on GPUs.
+
+---
+
+## Requirements
+
+- Python 3.8+
+- A GPU-enabled environment for CUDA support (optional but recommended).
+- Installed Python libraries:
+  - `gradio`
+  - `torch`
+  - `transformers`
+  - `torchvision`
+  - `Pillow`
+  - `numpy`
+
+Install dependencies using:
+```bash
+pip install gradio torch torchvision transformers Pillow numpy
+```
+
+---
+
+## Usage
+
+### Run the Application
+Execute the script using:
+```bash
+python inference.py
+```
+
+### Interface
+
+#### Tab 1: Image Upload
+1. Upload an image from your local system.
+2. Select a model (`BiRefNet` or `RMBG-2.0`).
+3. View and download the processed image with the background removed.
+
+#### Tab 2: URL Input
+1. Paste the URL of an image.
+2. Select a model (`BiRefNet` or `RMBG-2.0`).
+3. View and download the processed image with the background removed.
+
+#### Tab 3: File Output
+1. Upload an image file.
+2. Select a model (`BiRefNet` or `RMBG-2.0`).
+3. Get the path to the processed PNG file with transparency.
+
+### Example
+- Use the provided example image (`ironman.jpg`) to test the tool.
+
+
+---
+
+## How It Works
+
+1. **Model Loading**:
+   - Loads pre-trained segmentation models from Hugging Face.
+2. **Image Preprocessing**:
+   - Resizes and normalizes the input image.
+3. **Background Removal**:
+   - The model generates a mask for the image background.
+   - The mask is applied to create a transparent background.
+4. **Output**:
+   - Processed image is displayed or saved with an alpha channel.
+
+---
+
+
+## Contributing
+Feel free to submit issues or pull requests for improvements or bug fixes.
+
+---
+

inference.py

@@ -0,0 +1,173 @@
+
+import gradio as gr
+from load_image import load_img
+import spaces
+from transformers import AutoModelForImageSegmentation
+import torch
+from torchvision import transforms
+from PIL import Image
+import os
+import numpy as np
+
+torch.set_float32_matmul_precision(["high", "highest"][0])
+
+# load 2 models
+
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+).to("cuda")
+
+
+RMBG2 = AutoModelForImageSegmentation.from_pretrained(
+    "briaai/RMBG-2.0", trust_remote_code=True
+).to("cuda")
+
+# Keep them in a dict to switch easily
+models_dict = {
+    "BiRefNet": birefnet,
+    "RMBG-2.0": RMBG2,
+}
+
+# Transform
+
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
+
+@spaces.GPU
+def process(image: Image.Image, model_choice: str):
+    """
+    Runs inference to remove the background (adds alpha) 
+    with the chosen segmentation model.
+    """
+    # Select the model
+    current_model = models_dict[model_choice]
+
+    # Prepare image
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to("cuda")
+
+    # Inference
+    with torch.no_grad():
+        # Each model returns a list of preds in its forward, 
+        # so we take the last element, apply sigmoid, and move to CPU
+        preds = current_model(input_images)[-1].sigmoid().cpu()
+
+    # Convert single-channel pred to a PIL mask
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+
+    # Resize the mask back to original image size
+    mask = pred_pil.resize(image_size)
+
+    # Add alpha channel to the original
+    image.putalpha(mask)
+    return image
+
+def fn(source: str, model_choice: str):
+    """
+    Used by Tab 1 & Tab 2 to produce a processed image with alpha.
+    - 'source' is either a file path (type="filepath") or 
+      a URL string (textbox).
+    - 'model_choice' is the user's selection from the radio.
+    """
+    # Load from local path or URL
+    im = load_img(source, output_type="pil")
+    im = im.convert("RGB")
+
+    # Process
+    processed_image = process(im, model_choice)
+    return processed_image
+
+def process_file(file_path: str, model_choice: str):
+    """
+    For Tab 3 (file output).
+    - Accepts a local path, returns path to a new .png with alpha channel.
+    - 'model_choice' is also passed in for selecting the model.
+    """
+    name_path = file_path.rsplit(".", 1)[0] + ".png"
+    im = load_img(file_path, output_type="pil")
+    im = im.convert("RGB")
+
+    # Run the chosen model
+    transparent = process(im, model_choice)
+    transparent.save(name_path)
+    return name_path
+
+
+# GRadio UI
+
+model_selector_1 = gr.Radio(
+    choices=["BiRefNet", "RMBG-2.0"],
+    value="BiRefNet",
+    label="Select Model"
+)
+model_selector_2 = gr.Radio(
+    choices=["BiRefNet", "RMBG-2.0"],
+    value="BiRefNet",
+    label="Select Model"
+)
+model_selector_3 = gr.Radio(
+    choices=["BiRefNet", "RMBG-2.0"],
+    value="BiRefNet",
+    label="Select Model"
+)
+
+# Outputs for tabs 1 & 2: single processed image
+processed_img_upload = gr.Image(label="Processed Image (Upload)", type="pil")
+processed_img_url = gr.Image(label="Processed Image (URL)", type="pil")
+
+# For uploading local files
+image_upload = gr.Image(label="Upload an image", type="filepath")
+image_file_upload = gr.Image(label="Upload an image", type="filepath")
+
+# For Tab 2 (URL input)
+url_input = gr.Textbox(label="Paste an image URL")
+
+# For Tab 3 (file output)
+output_file = gr.File(label="Output PNG File")
+
+# Tab 1: local image -> processed image
+tab1 = gr.Interface(
+    fn=fn,
+    inputs=[image_upload, model_selector_1],
+    outputs=processed_img_upload,
+    examples=[["ironman.jpg", "BiRefNet/RMBG"]],
+    api_name="image",
+    description="Upload an image and choose your background removal model."
+)
+
+# Tab 2: URL input -> processed image
+tab2 = gr.Interface(
+    fn=fn,
+    inputs=[url_input, model_selector_2],
+    outputs=processed_img_url,
+    api_name="text",
+    description="Paste an image URL and choose your background removal model."
+)
+
+# Tab 3: file output -> returns path to .png
+tab3 = gr.Interface(
+    fn=process_file,
+    inputs=[image_file_upload, model_selector_3],
+    outputs=output_file,
+    examples=[["ironman.jpg", "BiRefNet/RMBG"]],
+    api_name="png",
+    description="Upload an image, choose a model, and get a transparent PNG."
+)
+
+# Combine all tabs
+demo = gr.TabbedInterface(
+    [tab1, tab2, tab3],
+    ["Image Upload", "URL Input", "File Output"],
+    title="Background Removal Tool"
+)
+
+if __name__ == "__main__":
+    demo.launch(show_error=True, share=True)
+
+

load_image.py

@@ -0,0 +1,35 @@
+
+import os
+import requests
+from io import BytesIO
+from PIL import Image
+import numpy as np
+
+def load_img(source, output_type="pil"):
+    """
+    Load an image from a local file path or a URL.
+
+    Parameters:
+    - source (str): A file path or a URL.
+    - output_type (str): The output format: "pil" (PIL Image) or "numpy" (NumPy array).
+
+    Returns:
+    - PIL.Image.Image or numpy.ndarray depending on output_type.
+    """
+
+    # Determine if `source` is a local file path or a URL
+    if os.path.exists(source):
+        # Local file
+        img = Image.open(source)
+    else:
+        # Assume source is a URL
+        response = requests.get(source)
+        response.raise_for_status()
+        img = Image.open(BytesIO(response.content))
+
+    if output_type == "pil":
+        return img
+    elif output_type == "numpy":
+        return np.array(img)
+    else:
+        raise ValueError(f"Unknown output_type: {output_type}")