Spaces:

videopix
/

background-remover-api

Running

App Files Files

videopix commited on Oct 15, 2025

Commit

e6e54a0

verified ·

1 Parent(s): 8ba660f

Update app.py

Browse files

Files changed (1) hide show

app.py +110 -27

app.py CHANGED Viewed

@@ -1,11 +1,15 @@
-import gradio as gr
 from PIL import Image
 from io import BytesIO
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
 from loadimg import load_img
-import os
 # -------------------------
 # Model Setup
@@ -13,16 +17,31 @@ import os
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print("Loading BiRefNet model...")
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet",
-    cache_dir=MODEL_DIR,
-    trust_remote_code=True,
-    revision="main"
-)
-birefnet.to(device).eval()
-print("Model loaded successfully.")
 # -------------------------
 # Image Preprocessing
@@ -33,7 +52,7 @@ def transform_image(image: Image.Image) -> torch.Tensor:
     mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
     std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
     arr = (arr - mean) / std
-    arr = np.transpose(arr, (2, 0, 1))
     tensor = torch.from_numpy(arr).unsqueeze(0).to(torch.float32).to(device)
     return tensor
@@ -49,22 +68,86 @@ def process_image(image: Image.Image) -> Image.Image:
     return image
 # -------------------------
-# Gradio Function
 # -------------------------
-def remove_background_gradio(input_img):
-    # Gradio passes PIL images directly
-    result = process_image(input_img.convert("RGB"))
-    return result
 # -------------------------
-# Gradio Interface
 # -------------------------
-iface = gr.Interface(
-    fn=remove_background_gradio,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Background Removal Tool",
-    description="Upload an image and get a transparent background."
-)
-iface.launch(server_name="0.0.0.0", server_port=7860)

+import os
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse
 from PIL import Image
 from io import BytesIO
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
 from loadimg import load_img
+from contextlib import asynccontextmanager
+import asyncio
+from functools import partial
 # -------------------------
 # Model Setup
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
+birefnet = None  # will initialize on startup
+# -------------------------
+# Lifespan (Startup/Shutdown)
+# -------------------------
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global birefnet
+    if birefnet is None:
+        print("Loading BiRefNet model...")
+        birefnet = AutoModelForImageSegmentation.from_pretrained(
+            "ZhengPeng7/BiRefNet",
+            cache_dir=MODEL_DIR,
+            trust_remote_code=True,
+            revision="main"
+        )
+        birefnet.to(device).eval()
+        print("Model loaded successfully.")
+    yield
+    # optional shutdown logic
+# -------------------------
+# FastAPI App
+# -------------------------
+app = FastAPI(title="Background Removal API", lifespan=lifespan)
 # -------------------------
 # Image Preprocessing
     mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
     std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
     arr = (arr - mean) / std
+    arr = np.transpose(arr, (2, 0, 1))  # HWC -> CHW
     tensor = torch.from_numpy(arr).unsqueeze(0).to(torch.float32).to(device)
     return tensor
     return image
 # -------------------------
+# /remove-background endpoint
 # -------------------------
+@app.post("/remove-background")
+async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
+    """
+    Accept either an uploaded file or an image URL.
+    Returns PNG with transparent background.
+    """
+    try:
+        if file:
+            image = Image.open(BytesIO(await file.read())).convert("RGB")
+        elif image_url:
+            image = load_img(image_url, output_type="pil").convert("RGB")
+        else:
+            raise HTTPException(status_code=400, detail="Provide file or image_url")
+        # run blocking image processing in a separate thread
+        loop = asyncio.get_running_loop()
+        result = await loop.run_in_executor(None, partial(process_image, image))
+        buf = BytesIO()
+        result.save(buf, format="PNG")
+        buf.seek(0)
+        return StreamingResponse(buf, media_type="image/png")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 # -------------------------
+# Optional: Web Interface
 # -------------------------
+@app.get("/", response_class=HTMLResponse)
+async def index():
+    html_content = """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Background Removal</title>
+        <style>
+            body { font-family: Arial; padding: 20px; }
+            .container { max-width: 600px; margin: auto; background: #f9f9f9; padding: 20px; border-radius: 10px; }
+            img { max-width: 100%; margin-top: 20px; }
+        </style>
+    </head>
+    <body>
+        <div class="container">
+            <h2>Background Removal Tool</h2>
+            <form id="fileForm" enctype="multipart/form-data">
+                <input type="file" name="file" id="fileInput">
+                <button type="submit">Remove Background</button>
+            </form>
+            <hr>
+            <form id="urlForm">
+                <input type="text" id="urlInput" placeholder="Image URL">
+                <button type="submit">Remove Background</button>
+            </form>
+            <img id="resultImg" src="">
+        </div>
+        <script>
+            const fileForm = document.getElementById('fileForm');
+            fileForm.addEventListener('submit', async e => {
+                e.preventDefault();
+                const fileInput = document.getElementById('fileInput');
+                if(fileInput.files.length === 0) return alert("Select a file!");
+                const formData = new FormData();
+                formData.append("file", fileInput.files[0]);
+                const res = await fetch('/remove-background', {method:'POST', body:formData});
+                const blob = await res.blob();
+                document.getElementById('resultImg').src = URL.createObjectURL(blob);
+            });
+            const urlForm = document.getElementById('urlForm');
+            urlForm.addEventListener('submit', async e => {
+                e.preventDefault();
+                const formData = new FormData();
+                formData.append("image_url", document.getElementById('urlInput').value);
+                const res = await fetch('/remove-background', {method:'POST', body:formData});
+                const blob = await res.blob();
+                document.getElementById('resultImg').src = URL.createObjectURL(blob);
+            });
+        </script>
+    </body>
+    </html>
+    """
+    return HTMLResponse(html_content)