Update app.py

app.py

@@ -1,12 +1,13 @@
 import os
-from io import BytesIO
-from PIL import Image, UnidentifiedImageError
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.responses import StreamingResponse
+from PIL import Image
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
-from fastapi import FastAPI, File, Form, UploadFile, HTTPException
-from fastapi.responses import StreamingResponse
+from io import BytesIO
 import requests
+import uvicorn
 
 # -------------------------
 # Model Setup
@@ -26,71 +27,71 @@ birefnet.to(device).eval()
 print("Model loaded successfully.")
 
 # -------------------------
-# Image Preprocessing
+# FastAPI App
+# -------------------------
+app = FastAPI(title="Background Remover API")
+
+# -------------------------
+# Utility Functions
 # -------------------------
+def load_image_from_url(url: str) -> Image.Image:
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        return Image.open(BytesIO(response.content)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error loading image from URL: {str(e)}")
+
 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))
+    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:
+    image_size = image.size
     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)
+        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.convert("RGBA")
     image.putalpha(mask)
     return image
 
 # -------------------------
-# FastAPI App
-# -------------------------
-app = FastAPI(title="Background Removal API")
-
-# -------------------------
-# API Endpoint: Return PNG
+# /remove-background Endpoint
 # -------------------------
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
     """
-    Accept either an uploaded file or an image URL.
-    Returns PNG with transparent background.
+    Remove background from an image.
+    Accepts either a file upload or an image URL.
+    Returns a PNG with transparent background.
     """
     try:
         if file:
-            img_bytes = await file.read()
-            img = Image.open(BytesIO(img_bytes)).convert("RGB")
+            image = Image.open(BytesIO(await file.read())).convert("RGB")
         elif image_url:
-            resp = requests.get(image_url, timeout=10)
-            resp.raise_for_status()
-            img = Image.open(BytesIO(resp.content)).convert("RGB")
+            image = load_image_from_url(image_url)
         else:
-            raise HTTPException(status_code=400, detail="Provide file or image_url")
+            raise HTTPException(status_code=400, detail="Provide either 'file' or 'image_url'.")
 
-        # Process background removal
-        result_img = process_image(img)
+        result = process_image(image)
 
-        # Convert to PNG bytes
         buf = BytesIO()
-        result_img.save(buf, format="PNG")
+        result.save(buf, format="PNG")
         buf.seek(0)
         return StreamingResponse(buf, media_type="image/png")
-    
-    except UnidentifiedImageError:
-        raise HTTPException(status_code=400, detail="Invalid image format")
-    except requests.RequestException:
-        raise HTTPException(status_code=400, detail="Failed to fetch image from URL")
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 # -------------------------
-# Optional Root Info
+# Run App on Spaces
 # -------------------------
-@app.get("/")
-async def index():
-    return {"message": "POST /remove-background with 'file' or 'image_url'. Returns PNG."}
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)