Update app.py

app.py

@@ -1,13 +1,14 @@
-import gradio as gr
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
-from PIL import Image
+import os
+import base64
+import requests
 from io import BytesIO
+from PIL import Image
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
-import os
-import requests
+import gradio as gr
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse, HTMLResponse
 
 # -------------------------
 # Model Setup
@@ -16,6 +17,7 @@ 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,
@@ -23,47 +25,61 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
     revision="main"
 )
 birefnet.to(device).eval()
+print("Model loaded successfully.")
 
-def transform_image(image: Image.Image):
+# -------------------------
+# 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])
-    std = np.array([0.229, 0.224, 0.225])
+    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))
-    return torch.from_numpy(arr).unsqueeze(0).to(torch.float32).to(device)
+    tensor = torch.from_numpy(arr).unsqueeze(0).to(torch.float32).to(device)
+    return tensor
 
-def process_image(image: Image.Image):
+def process_image(image: Image.Image) -> Image.Image:
     input_tensor = transform_image(image)
     with torch.no_grad():
-        pred = birefnet(input_tensor)[-1].sigmoid().cpu()[0,0]
-    mask = Image.fromarray((pred.numpy()*255).astype(np.uint8)).resize(image.size)
+        pred = birefnet(input_tensor)[-1].sigmoid().cpu()[0, 0]
+    mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(image.size)
+    image = image.convert("RGBA")
     image.putalpha(mask)
     return image
 
+# -------------------------
+# Gradio Function
+# -------------------------
 def remove_background_gradio(img):
     return process_image(img.convert("RGB"))
 
 # -------------------------
-# Gradio Interface
+# Gradio Interface for UI
 # -------------------------
 demo = 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."
 )
 
 # -------------------------
 # FastAPI App
 # -------------------------
-app = gr.routes.FastAPI.create_app(demo)  # Wraps Gradio app
+app = gr.routes.FastAPI.create_app(demo)  # Wrap Gradio
 
-# 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")
+    """
+    Custom endpoint: accepts 'file' upload or 'image_url' form.
+    Returns PNG bytes.
+    """
+    form = await request.form()
+    file = form.get("file")
+    image_url = form.get("image_url")
 
     if file:
         img = Image.open(file.file).convert("RGB")
@@ -71,10 +87,17 @@ async def remove_background(request: Request):
         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)
+        return {"error": "Provide file or image_url"}
 
-    result = remove_background_gradio(img)
+    result = process_image(img)
     buf = BytesIO()
     result.save(buf, format="PNG")
     buf.seek(0)
-    return JSONResponse({"image": "data:image/png;base64," + base64.b64encode(buf.read()).decode()})
+    return StreamingResponse(buf, media_type="image/png")
+
+# -------------------------
+# Optional: Root UI
+# -------------------------
+@app.get("/", response_class=HTMLResponse)
+async def index():
+    return demo.launch(share=False, inline=True)[0]  # Embed Gradio UI