Update app.py

app.py

@@ -1,80 +1,161 @@
-import io
 import os
-import threading
-from fastapi import FastAPI, File, UploadFile, Request
+from typing import Union
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import StreamingResponse, HTMLResponse
-from fastapi.templating import Jinja2Templates
 from PIL import Image
-import torch
-import torchvision.transforms as transforms
-import onnx
-import onnxruntime as ort
 import numpy as np
+import torch
+from transformers import AutoModelForImageSegmentation
+from io import BytesIO
+from loadimg import load_img
+import uvicorn
 
-# Settings
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-ONNX_PATH = os.path.join(os.path.dirname(__file__), "birefnet.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 safely
-if not os.path.exists(ONNX_PATH):
-    raise FileNotFoundError(f"ONNX model not found at {ONNX_PATH}")
-
-# Attempt to load external data automatically
-try:
-    onnx_session = ort.InferenceSession(ONNX_PATH, providers=["CUDAExecutionProvider"] if DEVICE=="cuda" else ["CPUExecutionProvider"])
-except ort.OnnxRuntimeError:
-    # Embed external data into memory if original session fails
-    print("Embedding external tensor data into the ONNX model...")
-    model = onnx.load(ONNX_PATH, load_external_data=False)  # embed data
-    embedded_path = ONNX_PATH.replace(".onnx", "_embedded.onnx")
-    onnx.save(model, embedded_path)
-    onnx_session = ort.InferenceSession(embedded_path, providers=["CUDAExecutionProvider"] if DEVICE=="cuda" else ["CPUExecutionProvider"])
+# -------------------------
+# Model Setup (Load Once)
+# -------------------------
+MODEL_DIR = "models/BiRefNet"
+os.makedirs(MODEL_DIR, exist_ok=True)
 
-print(f"ONNX model loaded with providers: {onnx_session.get_providers()}")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Lock for thread-safe inference
-onnx_lock = threading.Lock()
+print("Loading BiRefNet model (this may take a while on first run)...")
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet",
+    cache_dir=MODEL_DIR,
+    trust_remote_code=True
+)
+birefnet.to(device)
+birefnet.eval()
+print("Model loaded successfully.")
 
-def run_model_onnx(input_tensor: torch.Tensor) -> torch.Tensor:
-    with onnx_lock:
-        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
+# -------------------------
+# Image Preprocessing
+# -------------------------
+def transform_image(image: Image.Image) -> torch.Tensor:
+    image = image.resize((1024, 1024))
+    arr = np.array(image).astype(np.float32) / 255.0
+    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
 
 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
+    image_size = image.size
+    input_tensor = transform_image(image)
+    with torch.no_grad():
+        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+    pred = preds[0, 0]
+    mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(image_size)
+    image = image.copy()
+    image.putalpha(mask)
+    return image
 
-# FastAPI app
+# -------------------------
+# FastAPI App
+# -------------------------
 app = FastAPI(title="Background Removal API")
-templates = Jinja2Templates(directory="templates")
+
+# -------------------------
+# API Endpoints
+# -------------------------
+@app.post("/remove-background")
+async def remove_bg_file(file: UploadFile = File(...)):
+    """Upload an image file and get transparent PNG"""
+    try:
+        image = Image.open(BytesIO(await file.read())).convert("RGB")
+        transparent = process_image(image)
+        buf = BytesIO()
+        transparent.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))
 
 @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")
+async def remove_bg_url(image_url: str = Form(...)):
+    """Provide image URL and get transparent PNG"""
+    try:
+        image = load_img(image_url, output_type="pil").convert("RGB")
+        transparent = process_image(image)
+        buf = BytesIO()
+        transparent.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))
 
+# -------------------------
+# Web Interface
+# -------------------------
 @app.get("/", response_class=HTMLResponse)
-async def home(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request})
+async def index():
+    html_content = """
+    <!DOCTYPE html>
+    <html lang="en">
+    <head>
+        <meta charset="UTF-8">
+        <title>Background Removal Tool</title>
+        <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">
+        <style>
+            body { padding: 30px; background-color: #f8f9fa; }
+            .container { max-width: 600px; background: #fff; padding: 20px; border-radius: 10px; box-shadow: 0 0 10px rgba(0,0,0,0.1);}
+            img { max-width: 100%; margin-top: 10px; }
+        </style>
+    </head>
+    <body>
+        <div class="container">
+            <h2 class="mb-4">Background Removal Tool</h2>
+            <form id="fileForm" enctype="multipart/form-data">
+                <div class="mb-3">
+                    <label for="fileInput" class="form-label">Upload Image</label>
+                    <input class="form-control" type="file" id="fileInput" name="file">
+                </div>
+                <button class="btn btn-primary" type="submit">Remove Background</button>
+            </form>
+            <hr>
+            <form id="urlForm">
+                <div class="mb-3">
+                    <label for="urlInput" class="form-label">Image URL</label>
+                    <input class="form-control" type="text" id="urlInput" placeholder="Enter image URL">
+                </div>
+                <button class="btn btn-success" type="submit">Remove Background</button>
+            </form>
+            <hr>
+            <h5>Result:</h5>
+            <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_bg_file', {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 urlInput = document.getElementById('urlInput').value;
+                const formData = new FormData();
+                formData.append("image_url", urlInput);
+                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(content=html_content)
+
+# -------------------------
+# Run the server
+# -------------------------
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)