Update app.py

app.py

@@ -1,6 +1,5 @@
 import os
 from io import BytesIO
-import uvicorn
 
 import numpy as np
 import requests
@@ -11,19 +10,19 @@ from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 from transformers import AutoModelForImageSegmentation
 
-# -------------------------
-# Optional HEIC/HEIF Support
-# -------------------------
+# ---------------------------------------------------------
+# Optional HEIC Support
+# ---------------------------------------------------------
 try:
     import pillow_heif
     pillow_heif.register_heif_opener()
     print("HEIC/HEIF supported")
-except ImportError:
+except:
     print("Install pillow-heif for HEIC support")
 
-# -------------------------
-# Model Setup
-# -------------------------
+# ---------------------------------------------------------
+# Load Model
+# ---------------------------------------------------------
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 
@@ -38,130 +37,183 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
     revision="main"
 )
 birefnet.to(device, dtype=dtype).eval()
-print("Model loaded successfully.")
+print("Model Ready")
 
-# -------------------------
+# ---------------------------------------------------------
 # FastAPI App
-# -------------------------
+# ---------------------------------------------------------
 app = FastAPI(title="Background Remover API")
 
-# Allow API calls from mobile apps, web apps, backend servers
+# Allow requests from any app
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# -------------------------
+# ---------------------------------------------------------
 # Utility Functions
-# -------------------------
+# ---------------------------------------------------------
 def load_image_from_url(url: str):
     try:
         resp = requests.get(url, timeout=10)
         resp.raise_for_status()
         return Image.open(BytesIO(resp.content)).convert("RGB")
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Invalid image URL: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Error loading image URL: {str(e)}")
 
 def transform_image(image: Image.Image, resolution: int):
     image = image.resize((resolution, resolution))
-    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 = arr.transpose((2, 0, 1))
-    return torch.from_numpy(arr).unsqueeze(0).to(device=device, dtype=dtype)
+    img = np.array(image).astype("float32") / 255.0
+    mean = np.array([0.485, 0.456, 0.406], dtype="float32")
+    std = np.array([0.229, 0.224, 0.225], dtype="float32")
+    img = (img - mean) / std
+    img = img.transpose((2, 0, 1))
+    return torch.from_numpy(img).unsqueeze(0).to(device=device, dtype=dtype)
 
 def process_image(image: Image.Image, resolution: int):
-    orig = image.size
-    t = transform_image(image, resolution)
+    original_size = image.size
+    tensor = transform_image(image, resolution)
+
     with torch.no_grad():
-        pred = birefnet(t)[-1].sigmoid().cpu()[0, 0]
+        mask_pred = birefnet(tensor)[-1].sigmoid().cpu()[0, 0]
+
+    mask = Image.fromarray((mask_pred.numpy() * 255).astype("uint8"))
+    mask = mask.resize(original_size)
 
-    mask = Image.fromarray((pred.numpy() * 255).astype("uint8")).resize(orig)
     image = image.convert("RGBA")
     image.putalpha(mask)
     return image
 
-
-# -------------------------
-# GET + POST API SUPPORT
-# -------------------------
+# ---------------------------------------------------------
+# GET + POST endpoint
+# ---------------------------------------------------------
 @app.api_route("/remove-background", methods=["GET", "POST"])
 async def remove_background(
     file: UploadFile = File(None),
     image_url: str = Form(None),
     resolution: int = Form(512),
-    get_url: str = Query(None, description="Use for GET request: ?get_url=https://..."),
-    get_res: int = Query(512, description="Resolution for GET request"),
+    get_url: str = Query(None),
+    get_res: int = Query(512),
 ):
-    """
-    Supports:
-    - POST file upload
-    - POST image_url
-    - GET request using: /remove-background?get_url=...&get_res=512
-    """
-
     try:
-        # Determine GET or POST mode
+        # GET mode: /remove-background?get_url=...&get_res=512
         if get_url:
-            image = load_image_from_url(get_url)
+            img = load_image_from_url(get_url)
             resolution = get_res
 
+        # POST mode - file upload
         elif file:
-            content = await file.read()
-            if not content:
-                raise HTTPException(status_code=400, detail="Uploaded file is empty.")
-            image = Image.open(BytesIO(content)).convert("RGB")
+            data = await file.read()
+            if not data:
+                raise HTTPException(status_code=400, detail="Empty file")
+            img = Image.open(BytesIO(data)).convert("RGB")
 
+        # POST mode - URL
         elif image_url:
-            image = load_image_from_url(image_url)
+            img = load_image_from_url(image_url)
 
         else:
-            raise HTTPException(status_code=400, detail="No image provided.")
+            raise HTTPException(status_code=400, detail="No image provided")
 
-        result = process_image(image, resolution)
-        buf = BytesIO()
-        result.save(buf, format="PNG")
-        buf.seek(0)
+        result = process_image(img, resolution)
+
+        buffer = BytesIO()
+        result.save(buffer, format="PNG")
+        buffer.seek(0)
 
         return StreamingResponse(
-            buf,
+            buffer,
             media_type="image/png",
-            headers={"Content-Disposition": "inline; filename=result.png"}
+            headers={"Content-Disposition": "inline; filename=result.png"},
         )
 
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}")
 
-
-# -------------------------
-# Favicon handler (fix 404)
-# -------------------------
+# ---------------------------------------------------------
+# Favicon (stop 404 logs)
+# ---------------------------------------------------------
 @app.get("/favicon.ico")
 async def favicon():
     return HTMLResponse("")
 
-
-# -------------------------
-# Test Page
-# -------------------------
+# ---------------------------------------------------------
+# UI for testing (POST method)
+# ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 async def index():
     return """
-    <h2>Background Remover API Live</h2>
-    <p>POST endpoint: <code>/remove-background</code></p>
-    <p>GET example:</p>
-    <pre>/remove-background?get_url=https://example.com/img.jpg&get_res=512</pre>
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Background Remover API Tester</title>
+        <style>
+            body { font-family: Arial; padding: 20px; }
+            img { max-width: 100%; margin-top: 20px; border-radius: 10px; }
+        </style>
+    </head>
+    <body>
+        <h2>Background Remover API Test (POST)</h2>
+
+        <h4>Upload Image</h4>
+        <form id="uploadForm" enctype="multipart/form-data">
+            <input type="file" id="file" name="file"><br><br>
+            <label>Resolution:</label>
+            <input type="number" id="resFile" value="512"><br><br>
+            <button type="submit">Remove Background</button>
+        </form>
+
+        <h4>Or use Image URL</h4>
+        <form id="urlForm">
+            <input type="text" id="imgUrl" placeholder="https://example.com/image.jpg" size="50"><br><br>
+            <label>Resolution:</label>
+            <input type="number" id="resUrl" value="512"><br><br>
+            <button type="submit">Remove Background</button>
+        </form>
+
+        <h3>Result:</h3>
+        <img id="result" />
+
+        <script>
+            const resultImg = document.getElementById("result");
+
+            document.getElementById("uploadForm").onsubmit = async (e) => {
+                e.preventDefault();
+                const file = document.getElementById("file").files[0];
+                const res = document.getElementById("resFile").value;
+
+                const fd = new FormData();
+                fd.append("file", file);
+                fd.append("resolution", res);
+
+                const r = await fetch("/remove-background", { method: "POST", body: fd });
+                resultImg.src = URL.createObjectURL(await r.blob());
+            };
+
+            document.getElementById("urlForm").onsubmit = async (e) => {
+                e.preventDefault();
+                const url = document.getElementById("imgUrl").value;
+                const res = document.getElementById("resUrl").value;
+
+                const fd = new FormData();
+                fd.append("image_url", url);
+                fd.append("resolution", res);
+
+                const r = await fetch("/remove-background", { method: "POST", body: fd });
+                resultImg.src = URL.createObjectURL(await r.blob());
+            };
+        </script>
+    </body>
+    </html>
     """
 
-
-
-# -------------------------
+# ---------------------------------------------------------
 # Run App
-# -------------------------
+# ---------------------------------------------------------
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
-
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)