Update app.py

app.py

@@ -2,8 +2,8 @@ import os
 import threading
 import torch
 import numpy as np
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
+from fastapi.responses import StreamingResponse, HTMLResponse, RedirectResponse, JSONResponse
 from PIL import Image
 from io import BytesIO
 import requests
@@ -20,7 +20,7 @@ except ImportError:
     pass
 
 # ---------------------------------------------------------
-# Performance settings for CPU (HF Spaces)
+# Performance settings for HF CPU
 # ---------------------------------------------------------
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["MKL_NUM_THREADS"] = "1"
@@ -29,8 +29,8 @@ torch.set_num_threads(1)
 # ---------------------------------------------------------
 # Constants
 # ---------------------------------------------------------
-TARGET_SIZE = (512, 512)   # Faster inference resolution
-MAX_SIDE = 3000            # Auto-downscale limit for large uploads
+TARGET_SIZE = (512, 512)      # Faster inference
+MAX_SIDE = 3000               # Auto-downscale for huge uploads
 
 # ---------------------------------------------------------
 # Load model
@@ -46,67 +46,65 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet",
     cache_dir=MODEL_DIR,
     trust_remote_code=True,
-    revision="main"
+    revision="main",
 )
 birefnet.to(device, dtype=dtype).eval()
 print("Model ready.")
 
-# Thread lock for CPU safety
-inference_lock = threading.Lock()
+lock = threading.Lock()
 
 # ---------------------------------------------------------
-# Helper functions
+# Helpers
 # ---------------------------------------------------------
 def load_image_from_url(url: str) -> Image.Image:
     try:
         r = requests.get(url, timeout=10)
         r.raise_for_status()
         return Image.open(BytesIO(r.content)).convert("RGB")
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Cannot load image from URL: {str(e)}")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid image URL")
 
 
-def auto_downscale(image: Image.Image) -> Image.Image:
-    w, h = image.size
+def auto_downscale(img: Image.Image) -> Image.Image:
+    w, h = img.size
     if max(w, h) <= MAX_SIDE:
-        return image
+        return img
 
     scale = MAX_SIDE / max(w, h)
     new_w = int(w * scale)
     new_h = int(h * scale)
 
     print(f"[INFO] Downscaling {w}×{h} → {new_w}×{new_h}")
-    return image.resize((new_w, new_h), Image.LANCZOS)
-
+    return img.resize((new_w, new_h), Image.LANCZOS)
 
-def transform_image(image: Image.Image) -> torch.Tensor:
-    image = image.resize(TARGET_SIZE)
 
-    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)
+def transform(img: Image.Image) -> torch.Tensor:
+    img = img.resize(TARGET_SIZE)
 
+    arr = np.array(img).astype(np.float32) / 255.0
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
     arr = (arr - mean) / std
     arr = np.transpose(arr, (2, 0, 1))
 
-    tensor = torch.from_numpy(arr).unsqueeze(0).to(device=device, dtype=dtype)
-    return tensor
+    t = torch.from_numpy(arr).unsqueeze(0).to(device=device, dtype=dtype)
+    return t
 
 
-def run_inference(image: Image.Image) -> Image.Image:
-    orig_size = image.size
-    input_tensor = transform_image(image)
+def run_inference(img: Image.Image) -> Image.Image:
+    orig_size = img.size
+    tensor = transform(img)
 
-    with inference_lock:
+    with lock:
         with torch.no_grad():
-            preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+            pred = birefnet(tensor)[-1].sigmoid().cpu()[0, 0]
 
-    pred = preds[0, 0]
     mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(orig_size)
 
-    image = image.convert("RGBA")
-    image.putalpha(mask)
-    return image
+    img = img.convert("RGBA")
+    img.putalpha(mask)
+    return img
+
 
 # ---------------------------------------------------------
 # FastAPI app
@@ -114,103 +112,122 @@ def run_inference(image: Image.Image) -> Image.Image:
 app = FastAPI(title="Background Remover API")
 
 # ---------------------------------------------------------
-# POST endpoint only (no GET processing)
+# Redirect GET → POST logic
+# ---------------------------------------------------------
+@app.get("/remove-background")
+async def redirect_to_post():
+    return JSONResponse(
+        {"detail": "This endpoint only supports POST. Use POST /remove-background"},
+        status_code=405
+    )
+
+# ---------------------------------------------------------
+# Main POST endpoint
 # ---------------------------------------------------------
 @app.post("/remove-background")
-async def remove_background(
-    file: UploadFile = File(None),
-    image_url: str = Form(None)
-):
+async def remove_bg(file: UploadFile = File(None), image_url: str = Form(None)):
     try:
-        # load image
         if file:
             raw = await file.read()
-            image = Image.open(BytesIO(raw)).convert("RGB")
+            img = Image.open(BytesIO(raw)).convert("RGB")
         elif image_url:
-            image = load_image_from_url(image_url)
+            img = load_image_from_url(image_url)
         else:
-            raise HTTPException(status_code=400, detail="Provide file or image_url.")
-
-        # auto shrink large inputs
-        image = auto_downscale(image)
+            raise HTTPException(status_code=400, detail="Upload file or image_url required")
 
-        # remove background
-        result = run_inference(image)
+        img = auto_downscale(img)
+        result = run_inference(img)
 
-        # return PNG
         buf = BytesIO()
-        result.save(buf, format="PNG", optimize=True)
+        result.save(buf, format="PNG")
         buf.seek(0)
 
         return StreamingResponse(buf, media_type="image/png")
 
-    except HTTPException:
-        raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
 
 # ---------------------------------------------------------
-# UI for POST method testing only
+# UI: Show INPUT + OUTPUT (big preview)
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 async def ui():
     return """
     <html>
     <head>
-        <title>Background Remover – Test Tool</title>
-        <link rel='stylesheet' href='https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css'>
+        <title>Background Remover – Test UI</title>
+        <link rel='stylesheet'
+              href='https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css'>
     </head>
-
-    <body style='background:#f8f9fa;padding-top:40px;'>
-        <div class='container text-center'>
-            <h2>POST Method Test Panel</h2>
-            <p>This UI only sends POST requests to <code>/remove-background</code>.</p>
-
-            <h5>Test with File Upload:</h5>
-            <form id='uploadForm' enctype='multipart/form-data'>
-                <input class='form-control mb-2' type='file' id='fileInput'>
+    <body class='bg-light'>
+        <div class='container py-4 text-center'>
+
+            <h2 class='mb-4'>API Test Panel (POST Only)</h2>
+
+            <div class='row'>
+                <div class='col-md-6'>
+                    <h5>Input Image</h5>
+                    <img id='inputImg' style='max-width:100%; border-radius:10px;'>
+                </div>
+                <div class='col-md-6'>
+                    <h5>Output Image</h5>
+                    <img id='outputImg' style='max-width:100%; border-radius:10px;'>
+                </div>
+            </div>
+
+            <hr>
+
+            <h4>Upload Test</h4>
+            <form id="uploadForm" enctype='multipart/form-data'>
+                <input type='file' id='fileInput' class='form-control mb-3'>
                 <button class='btn btn-primary'>Send POST</button>
             </form>
 
-            <div class='my-4'>OR</div>
+            <hr>
 
-            <h5>Test with Image URL:</h5>
+            <h4>URL Test</h4>
             <form id='urlForm'>
-                <input class='form-control mb-2' id='urlInput' placeholder='https://example.com/image.jpg'>
+                <input id='urlInput' class='form-control mb-3' placeholder='https://example.com/image.jpg'>
                 <button class='btn btn-success'>Send POST</button>
             </form>
-
-            <h4 class='mt-4'>Output:</h4>
-            <img id='outputImg' style='max-width:90%;border-radius:10px;'/>
         </div>
 
         <script>
-            const outputImg = document.getElementById("outputImg");
+        const inputImg = document.getElementById("inputImg");
+        const outputImg = document.getElementById("outputImg");
 
-            document.getElementById("uploadForm").addEventListener("submit", async e => {
-                e.preventDefault();
-                const file = document.getElementById("fileInput").files[0];
-                if (!file) return alert("Select a file first");
+        // FILE TEST
+        document.getElementById("uploadForm").addEventListener("submit", async e => {
+            e.preventDefault();
+            const file = document.getElementById("fileInput").files[0];
+            if (!file) return alert("Select a file first.");
 
-                const fd = new FormData();
-                fd.append("file", file);
+            inputImg.src = URL.createObjectURL(file);
 
-                const r = await fetch("/remove-background", {method:"POST", body:fd});
-                outputImg.src = URL.createObjectURL(await r.blob());
-            });
+            const fd = new FormData();
+            fd.append("file", file);
 
-            document.getElementById("urlForm").addEventListener("submit", async e => {
-                e.preventDefault();
-                const url = document.getElementById("urlInput").value.trim();
-                if (!url) return alert("Enter an image URL");
+            const r = await fetch("/remove-background", { method:"POST", body:fd });
+            outputImg.src = URL.createObjectURL(await r.blob());
+        });
 
-                const fd = new FormData();
-                fd.append("image_url", url);
+        // URL TEST
+        document.getElementById("urlForm").addEventListener("submit", async e => {
+            e.preventDefault();
+            const url = document.getElementById("urlInput").value.trim();
+            if (!url) return alert("Enter an image URL first.");
 
-                const r = await fetch("/remove-background", {method:"POST", body:fd});
-                outputImg.src = URL.createObjectURL(await r.blob());
-            });
+            inputImg.src = url;
+
+            const fd = new FormData();
+            fd.append("image_url", url);
+
+            const r = await fetch("/remove-background", { method:"POST", body:fd });
+            outputImg.src = URL.createObjectURL(await r.blob());
+        });
         </script>
+
     </body>
     </html>
     """