Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 from io import BytesIO
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
@@ -52,6 +52,9 @@ def process_image(image: Image.Image) -> Image.Image:
 # -------------------------
 app = FastAPI(title="Background Removal API")
 
+# -------------------------
+# API Endpoint: Return PNG
+# -------------------------
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
     """
@@ -60,17 +63,34 @@ async def remove_background(file: UploadFile = File(None), image_url: str = Form
     """
     try:
         if file:
-            img = Image.open(BytesIO(await file.read())).convert("RGB")
+            img_bytes = await file.read()
+            img = Image.open(BytesIO(img_bytes)).convert("RGB")
         elif image_url:
-            resp = requests.get(image_url)
+            resp = requests.get(image_url, timeout=10)
+            resp.raise_for_status()
             img = Image.open(BytesIO(resp.content)).convert("RGB")
         else:
             raise HTTPException(status_code=400, detail="Provide file or image_url")
 
-        result = process_image(img)
+        # Process background removal
+        result_img = process_image(img)
+
+        # Convert to PNG bytes
         buf = BytesIO()
-        result.save(buf, format="PNG")
+        result_img.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
+# -------------------------
+@app.get("/")
+async def index():
+    return {"message": "POST /remove-background with 'file' or 'image_url'. Returns PNG."}