Spaces:

videopix
/

background-remover-api

Running

App Files Files

videopix commited on Oct 15, 2025

Commit

c77bffa

verified ·

1 Parent(s): 347af97

Update app.py

Browse files

Files changed (1) hide show

app.py +124 -39

app.py CHANGED Viewed

@@ -1,77 +1,162 @@
 import io
 import os
-from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import StreamingResponse, HTMLResponse
-from PIL import Image, UnidentifiedImageError
 import torch
 import torchvision.transforms as transforms
 import onnxruntime as ort
-# -----------------------------
 # Settings
-# -----------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-ONNX_PATH = "BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx"
 if not os.path.exists(ONNX_PATH):
     raise FileNotFoundError(f"ONNX model not found at {ONNX_PATH}")
-# Use only one model, load it once
 providers = ["CUDAExecutionProvider"] if DEVICE == "cuda" else ["CPUExecutionProvider"]
 onnx_session = ort.InferenceSession(ONNX_PATH, providers=providers)
-# Preprocessing
-transform_image = transforms.Compose([
-    transforms.Resize((1024, 1024)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-])
-# -----------------------------
-# Helper
-# -----------------------------
-def run_model(image: Image.Image):
-    input_tensor = transform_image(image).unsqueeze(0)
-    ort_inputs = {onnx_session.get_inputs()[0].name: input_tensor.cpu().numpy()}
-    ort_outs = onnx_session.run(None, ort_inputs)
-    pred = torch.from_numpy(ort_outs[-1])[0]
     if pred.dim() == 3:
         pred = pred[0].squeeze(0)
     mask = transforms.ToPILImage()(pred.clamp(0, 1))
-    mask = mask.resize(image.size)
-    result = image.convert("RGBA")
-    result.putalpha(mask)
-    return result
-# -----------------------------
 # FastAPI app
-# -----------------------------
 app = FastAPI(title="Background Removal API")
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(...)):
-    try:
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-    except UnidentifiedImageError:
-        raise HTTPException(status_code=400, detail="Unsupported or corrupted image file.")
-    result_image = run_model(image)
     buf = io.BytesIO()
     result_image.save(buf, format="PNG")
     buf.seek(0)
     return StreamingResponse(buf, media_type="image/png")
 @app.get("/", response_class=HTMLResponse)
-async def home():
     return """
     <html>
     <body>
-        <h1>Background Remover</h1>
-        <form method="post" enctype="multipart/form-data" action="/remove-background">
-            <input type="file" name="file" accept="image/*" required>
-            <button type="submit">Remove Background</button>
-        </form>
     </body>
     </html>
-    """

 import io
 import os
+import threading
+from fastapi import FastAPI, File, UploadFile, Request
 from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from PIL import Image
 import torch
 import torchvision.transforms as transforms
 import onnxruntime as ort
 # Settings
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+ONNX_PATH = os.path.join(os.path.dirname(__file__), "BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx")
+# Preprocessing 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]),
+])
+# Load ONNX model
 if not os.path.exists(ONNX_PATH):
     raise FileNotFoundError(f"ONNX model not found at {ONNX_PATH}")
 providers = ["CUDAExecutionProvider"] if DEVICE == "cuda" else ["CPUExecutionProvider"]
 onnx_session = ort.InferenceSession(ONNX_PATH, providers=providers)
+print(f"ONNX model loaded with providers: {providers}")
+# Lock for thread-safe ONNX inference
+onnx_lock = threading.Lock()
+# Helper functions
+def run_model_onnx(input_tensor: torch.Tensor) -> torch.Tensor:
+    with onnx_lock:  # ensure thread safety
+        ort_inputs = {onnx_session.get_inputs()[0].name: input_tensor.cpu().numpy()}
+        ort_outs = onnx_session.run(None, ort_inputs)
+    preds = torch.from_numpy(ort_outs[-1]).sigmoid()
+    return preds
+def process_image(image: Image.Image) -> Image.Image:
+    original_size = image.size
+    input_tensor = transform_image(image).unsqueeze(0)  # (1,C,H,W)
+    preds = run_model_onnx(input_tensor)
+    pred = preds[0]
     if pred.dim() == 3:
         pred = pred[0].squeeze(0)
     mask = transforms.ToPILImage()(pred.clamp(0, 1))
+    mask = mask.resize(original_size, resample=Image.BILINEAR)
+    image_rgba = image.convert("RGBA")
+    image_rgba.putalpha(mask)
+    return image_rgba
 # FastAPI app
 app = FastAPI(title="Background Removal API")
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(...)):
+    image = Image.open(file.file).convert("RGB")
+    result_image = process_image(image)
     buf = io.BytesIO()
     result_image.save(buf, format="PNG")
     buf.seek(0)
     return StreamingResponse(buf, media_type="image/png")
+# Serve a simple HTML frontend for testing
 @app.get("/", response_class=HTMLResponse)
+async def home(request: Request):
     return """
+    <!DOCTYPE html>
     <html>
+    <head>
+        <title>Background Remover</title>
+        <style>
+            body {
+                font-family: Arial, sans-serif;
+                display: flex;
+                flex-direction: column;
+                align-items: center;
+                justify-content: center;
+                min-height: 100vh;
+                margin: 0;
+                padding: 20px;
+                background: #f5f5f5;
+            }
+            h1 {
+                color: #333;
+            }
+            .container {
+                background: white;
+                padding: 20px;
+                border-radius: 12px;
+                box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+                max-width: 500px;
+                width: 100%;
+                text-align: center;
+            }
+            input[type=file] {
+                margin: 15px 0;
+            }
+            button {
+                background: #4CAF50;
+                color: white;
+                padding: 10px 20px;
+                border: none;
+                border-radius: 6px;
+                cursor: pointer;
+                font-size: 16px;
+            }
+            button:hover {
+                background: #45a049;
+            }
+            img {
+                margin-top: 20px;
+                max-width: 100%;
+                border-radius: 8px;
+            }
+        </style>
+    </head>
     <body>
+        <div class="container">
+            <h1>Background Remover</h1>
+            <form id="upload-form">
+                <input type="file" id="file-input" name="file" accept="image/*" required />
+                <br/>
+                <button type="submit">Remove Background</button>
+            </form>
+            <div id="result"></div>
+        </div>
+        <script>
+            const form = document.getElementById('upload-form');
+            const resultDiv = document.getElementById('result');
+            form.addEventListener('submit', async (e) => {
+                e.preventDefault();
+                const fileInput = document.getElementById('file-input');
+                if (!fileInput.files.length) return;
+                const formData = new FormData();
+                formData.append('file', fileInput.files[0]);
+                resultDiv.innerHTML = "<p>Processing...</p>";
+                try {
+                    const response = await fetch('/remove-background', {
+                        method: 'POST',
+                        body: formData
+                    });
+                    if (!response.ok) {
+                        resultDiv.innerHTML = "<p style='color:red;'>Error processing image</p>";
+                        return;
+                    }
+                    const blob = await response.blob();
+                    const url = URL.createObjectURL(blob);
+                    resultDiv.innerHTML = `<h3>Result:</h3><img src="${url}" alt="Processed Image"/>`;
+                } catch (err) {
+                    resultDiv.innerHTML = "<p style='color:red;'>Request failed</p>";
+                }
+            });
+        </script>
     </body>
     </html>
+    """