Update app.py

app.py

@@ -1,13 +1,13 @@
-import os
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from fastapi.responses import StreamingResponse, HTMLResponse
+import gradio as gr
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
 from PIL import Image
+from io import BytesIO
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
-from io import BytesIO
-from loadimg import load_img
-from contextlib import asynccontextmanager
+import os
+import requests
 
 # -------------------------
 # Model Setup
@@ -15,134 +15,66 @@ from contextlib import asynccontextmanager
 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
-    # shutdown logic (optional)
-    
-# -------------------------
-# FastAPI App
-# -------------------------
-app = FastAPI(title="Background Removal API", lifespan=lifespan)
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet",
+    cache_dir=MODEL_DIR,
+    trust_remote_code=True,
+    revision="main"
+)
+birefnet.to(device).eval()
 
-# -------------------------
-# Image Processing
-# -------------------------
-def transform_image(image: Image.Image) -> torch.Tensor:
+def transform_image(image: Image.Image):
     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)
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
     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
+    arr = np.transpose(arr, (2, 0, 1))
+    return torch.from_numpy(arr).unsqueeze(0).to(torch.float32).to(device)
 
-def process_image(image: Image.Image) -> Image.Image:
-    image_size = image.size
+def process_image(image: Image.Image):
     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()
+        pred = birefnet(input_tensor)[-1].sigmoid().cpu()[0,0]
+    mask = Image.fromarray((pred.numpy()*255).astype(np.uint8)).resize(image.size)
     image.putalpha(mask)
     return image
 
+def remove_background_gradio(img):
+    return process_image(img.convert("RGB"))
+
 # -------------------------
-# API Endpoint
+# Gradio Interface
 # -------------------------
-@app.post("/remove-background")
-async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
-    """
-    Accept either an uploaded file or 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")
-        
-        result = 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))
+demo = gr.Interface(
+    fn=remove_background_gradio,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+)
 
 # -------------------------
-# Web Interface
+# FastAPI App
 # -------------------------
-@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)
+app = gr.routes.FastAPI.create_app(demo)  # Wraps Gradio app
+
+# Custom route for `/remove-background`
+@app.post("/remove-background")
+async def remove_background(request: Request):
+    data = await request.form()
+    file = data.get("file")
+    image_url = data.get("image_url")
+
+    if file:
+        img = Image.open(file.file).convert("RGB")
+    elif image_url:
+        resp = requests.get(image_url)
+        img = Image.open(BytesIO(resp.content)).convert("RGB")
+    else:
+        return JSONResponse({"error": "Provide file or image_url"}, status_code=400)
+
+    result = remove_background_gradio(img)
+    buf = BytesIO()
+    result.save(buf, format="PNG")
+    buf.seek(0)
+    return JSONResponse({"image": "data:image/png;base64," + base64.b64encode(buf.read()).decode()})