Update app.py

app.py

@@ -1,14 +1,12 @@
 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 gradio as gr
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi import FastAPI, File, Form, UploadFile, HTTPException
+from fastapi.responses import StreamingResponse
+import requests
 
 # -------------------------
 # Model Setup
@@ -49,55 +47,30 @@ def process_image(image: Image.Image) -> Image.Image:
     image.putalpha(mask)
     return image
 
-# -------------------------
-# Gradio Function
-# -------------------------
-def remove_background_gradio(img):
-    return process_image(img.convert("RGB"))
-
-# -------------------------
-# 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)  # Wrap Gradio
+app = FastAPI(title="Background Removal API")
 
 @app.post("/remove-background")
-async def remove_background(request: Request):
+async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
     """
-    Custom endpoint: accepts 'file' upload or 'image_url' form.
-    Returns PNG bytes.
+    Accept either an uploaded file or an image URL.
+    Returns PNG with transparent background.
     """
-    form = await request.form()
-    file = form.get("file")
-    image_url = form.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 {"error": "Provide file or image_url"}
+    try:
+        if file:
+            img = Image.open(BytesIO(await file.read())).convert("RGB")
+        elif image_url:
+            resp = requests.get(image_url)
+            img = Image.open(BytesIO(resp.content)).convert("RGB")
+        else:
+            raise HTTPException(status_code=400, detail="Provide file or image_url")
 
-    result = process_image(img)
-    buf = BytesIO()
-    result.save(buf, format="PNG")
-    buf.seek(0)
-    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
+        result = process_image(img)
+        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))