Initial commit - RMBG background remover FastAPI app

Dockerfile

@@ -0,0 +1,17 @@
+# Base image
+FROM python:3.10-slim
+
+# Set working directory
+WORKDIR /app
+
+# Copy files
+COPY . .
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Expose port
+EXPOSE 7860
+
+# Run FastAPI app
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,12 +1,12 @@
+# 🧠 Background Remover (RMBG-1.4)
+
+A FastAPI-based background-removal AI using the **BriaAI RMBG-1.4** model  
+(Improved U²-Net, deployed on Hugging Face Spaces).
+
 ---
-title: Removeit
-emoji: 🐠
-colorFrom: yellow
-colorTo: red
-sdk: docker
-pinned: false
-license: mit
-short_description: Ai powered background remover
----
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## 🚀 Run locally
+
+```bash
+pip install -r requirements.txt
+uvicorn app:app --reload

app.py

@@ -0,0 +1,16 @@
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import StreamingResponse
+from utils.rmbg_model import RMBGRemover
+
+app = FastAPI(title="Background Remover (RMBG-1.4)")
+remover = RMBGRemover()
+
+@app.get("/")
+def root():
+    return {"message": "RMBG-1.4 Background Remover is running 🚀"}
+
+@app.post("/remove-bg")
+async def remove_bg(file: UploadFile = File(...)):
+    image_bytes = await file.read()
+    output = remover.remove_background(image_bytes)
+    return StreamingResponse(output, media_type="image/png")

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+numpy
+pillow
+fastapi
+uvicorn
+python-multipart
+scikit-image
+transformers

utils/rmbg_model.py

@@ -0,0 +1,52 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+from transformers import AutoModelForImageSegmentation
+from PIL import Image
+from skimage import io
+import io as sysio
+
+
+class RMBGRemover:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = AutoModelForImageSegmentation.from_pretrained(
+            "briaai/RMBG-1.4", trust_remote_code=True
+        ).to(self.device)
+        self.model.eval()
+
+    def preprocess(self, im: np.ndarray, model_input_size: list) -> torch.Tensor:
+        if len(im.shape) < 3:
+            im = im[:, :, np.newaxis]
+        im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
+        im_tensor = F.interpolate(im_tensor.unsqueeze(0), size=model_input_size, mode="bilinear")
+        image = im_tensor / 255.0
+        image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+        return image
+
+    def postprocess(self, result: torch.Tensor, im_size: list) -> np.ndarray:
+        result = F.interpolate(result, size=im_size, mode="bilinear").squeeze(0)
+        result = (result - result.min()) / (result.max() - result.min())
+        im_array = (result * 255).permute(1, 2, 0).cpu().numpy().astype(np.uint8)
+        return np.squeeze(im_array)
+
+    def remove_background(self, image_bytes: bytes) -> bytes:
+        im = Image.open(sysio.BytesIO(image_bytes)).convert("RGB")
+        np_im = np.array(im)
+        im_size = np_im.shape[0:2]
+
+        model_input_size = [1024, 1024]
+        image = self.preprocess(np_im, model_input_size).to(self.device)
+
+        with torch.no_grad():
+            result = self.model(image)
+
+        mask = self.postprocess(result[0][0], im_size)
+        pil_mask = Image.fromarray(mask)
+        im.putalpha(pil_mask)
+
+        out_bytes = sysio.BytesIO()
+        im.save(out_bytes, format="PNG")
+        out_bytes.seek(0)
+        return out_bytes