Upload 17 files

README.md

@@ -1,41 +1,58 @@
----
-license: other
-license_name: custom-non-commercial
-license_link: https://huggingface.co/Muhammad8815/inno-rmbg-removal-8815/blob/main/LICENSE
-pipeline_tag: image-segmentation
-tags:
-- remove-background
-- background-removal
-- pytorch
-- onnx
-- image-segmentation
----
-
-# Inno RMBG Removal v1.0
-
-This is a custom-trained image background removal model using PyTorch and ONNX. It performs pixel-wise background removal on input images with high accuracy.
-
-## Model Highlights
-
-- Supports PyTorch `.pth` and ONNX inference.
-- Trained on diverse images of objects, people, and scenes.
-- Provides a binary mask of the foreground, suitable for downstream tasks (e.g. background replacement, transparency effects).
-
-## Example Input/Output
-
-**Input:**  
-![example input](example_input.jpg)
-
-**Output (no background):**  
-![result](results.png)
-
-## Inference Code (PyTorch)
-```python
-from PIL import Image
-from utilities import preprocess_image, postprocess_image, load_model
-
-model = load_model("model.pth")
-image = Image.open("example_input.jpg")
-mask = model.predict(image)
-image.putalpha(mask)
-image.save("output.png")
+---
+title: Inno Background Remover
+emoji: 🪄
+colorFrom: pink
+colorTo: purple
+sdk: gradio
+sdk_version: "4.19.2"
+app_file: app.py
+pinned: false
+---
+
+# 🪄 Inno RMBG Removal v1.0
+
+This Space runs a custom-trained image background removal model using PyTorch and ONNX.  
+It performs pixel-wise background removal on input images with high accuracy using a Gradio interface.
+
+---
+
+## ✨ Model Highlights
+
+- Trained on diverse images of objects, people, and scenes.
+- Outputs a binary mask of the foreground.
+- Works with `.pth` and ONNX formats.
+- Useful for background replacement or transparent cutouts.
+
+---
+
+## 📦 Input / Output
+
+**Input Image:**
+
+![example input](example_input.jpg)
+
+**Output (no background):**
+
+![result](results.png)
+
+---
+
+## 🧪 How It Works
+
+1. Upload an image using the UI.
+2. The model generates a mask for the foreground.
+3. The output is a transparent PNG or white background version.
+
+---
+
+## ⚙️ Inference Code (example)
+
+```python
+from PIL import Image
+from utilities import preprocess_image, postprocess_image, load_model
+
+model = load_model("model.pth")
+image = Image.open("example_input.jpg")
+mask = model.predict(image)
+image.putalpha(mask)
+image.save("output.png")

app.py

@@ -0,0 +1,67 @@
+import gradio as gr
+from briarmbg import BriaRMBG
+from PIL import Image
+import torch
+import torchvision.transforms as T
+import os
+import json
+from fastapi import Request, HTTPException
+
+# ===== Load API key from Hugging Face secret or fallback =====
+API_KEY = os.getenv("API_KEY", "sk_muhammad8815_Z4gXr93nPqT2LwV8")
+
+# ===== Load model =====
+model = BriaRMBG.from_pretrained("./")
+model.load_state_dict(torch.load("model.pth", map_location="cpu"))
+model.eval()
+
+# ===== Define preprocessing =====
+transform = T.Compose([
+    T.Resize((512, 512)),
+    T.ToTensor()
+])
+
+# ===== Background removal function with API key check =====
+def remove_background(image, request: Request = None):
+    # Check for API key in Authorization header
+    if request is not None:
+        auth = request.headers.get("authorization")
+        if not auth or not auth.startswith("Bearer ") or auth.split(" ")[1] != API_KEY:
+            raise HTTPException(status_code=401, detail="Invalid or missing API key")
+
+    # Preprocess image
+    img = transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        result = model(img)
+
+    # Handle different result formats
+    if isinstance(result, dict) and "pred" in result:
+        result = result["pred"]
+    elif isinstance(result, (tuple, list)):
+        result = result[0]
+    if isinstance(result, list):
+        result = result[0]
+
+    result = result.squeeze().numpy()
+
+    # Apply transparency mask
+    image = image.resize((result.shape[1], result.shape[0]))
+    image = image.convert("RGBA")
+    pixels = image.load()
+
+    for y in range(image.height):
+        for x in range(image.width):
+            if result[y][x] < 0.5:
+                pixels[x, y] = (255, 255, 255, 0)
+
+    return image
+
+# ===== Launch Gradio app with secure FastAPI request passthrough =====
+gr.Interface(
+    fn=remove_background,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Background Remover",
+    allow_flagging="never"
+).launch()

briarmbg.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PreTrainedModel
-from .MyConfig import RMBGConfig
+from MyConfig import RMBGConfig
 
 class REBNCONV(nn.Module):
     def __init__(self,in_ch=3,out_ch=3,dirate=1,stride=1):