Spaces:

videopix
/

background-remover-api

Running

App Files Files

videopix commited on Oct 15, 2025

Commit

4919185

verified ·

1 Parent(s): 1282773

Update app.py

Browse files

Files changed (1) hide show

app.py +29 -107

app.py CHANGED Viewed

@@ -1,47 +1,28 @@
-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 = "cpu"  # force CPU usage
-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 on CPU...")
-        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
@@ -52,7 +33,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))  # HWC -> CHW
     tensor = torch.from_numpy(arr).unsqueeze(0).to(torch.float32).to(device)
     return tensor
@@ -68,82 +49,23 @@ def process_image(image: Image.Image) -> Image.Image:
     return image
 # -------------------------
-# /remove-background endpoint
 # -------------------------
-@app.post("/remove-background")
-async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
-    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)

+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
 # -------------------------
 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
     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
     return image
 # -------------------------
+# Gradio Function for API
 # -------------------------
+def remove_background_gradio(input_img: Image.Image) -> Image.Image:
+    return process_image(input_img.convert("RGB"))
 # -------------------------
+# Gradio Interface (Web + API)
 # -------------------------
+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.",
+    allow_flagging="never",
+    api_name="remove-background"  # This exposes /api/predict/remove-background
+)
+# Launch
+iface.launch(server_name="0.0.0.0", server_port=7860)