Update app.py

app.py

@@ -11,7 +11,7 @@ from transformers import AutoModelForImageSegmentation
 import uvicorn
 
 # ---------------------------------------------------------
-# Optional HEIC/HEIF support
+# Optional HEIC/HEIF
 # ---------------------------------------------------------
 try:
     import pillow_heif
@@ -26,6 +26,12 @@ os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["MKL_NUM_THREADS"] = "1"
 torch.set_num_threads(1)
 
+# ---------------------------------------------------------
+# Constants
+# ---------------------------------------------------------
+TARGET_SIZE = (512, 512)   # Faster inference resolution
+MAX_SIDE = 3000            # Auto-downscale limit for large uploads
+
 # ---------------------------------------------------------
 # Load model
 # ---------------------------------------------------------
@@ -35,7 +41,7 @@ os.makedirs(MODEL_DIR, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-print("Loading model...")
+print("Loading BiRefNet…")
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet",
     cache_dir=MODEL_DIR,
@@ -45,7 +51,7 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
 birefnet.to(device, dtype=dtype).eval()
 print("Model ready.")
 
-# Thread lock for safe inference on CPU
+# Thread lock for CPU safety
 inference_lock = threading.Lock()
 
 # ---------------------------------------------------------
@@ -60,36 +66,36 @@ def load_image_from_url(url: str) -> Image.Image:
         raise HTTPException(status_code=400, detail=f"Cannot load image from URL: {str(e)}")
 
 
-def auto_downscale(image: Image.Image, max_side: int = 3000) -> Image.Image:
+def auto_downscale(image: Image.Image) -> Image.Image:
     w, h = image.size
-    if max(w, h) <= max_side:
+    if max(w, h) <= MAX_SIDE:
         return image
 
-    scale = max_side / max(w, h)
+    scale = MAX_SIDE / max(w, h)
     new_w = int(w * scale)
     new_h = int(h * scale)
 
-    print(f"[INFO] Downscaling large image {w}x{h} → {new_w}x{new_h}")
+    print(f"[INFO] Downscaling {w}×{h} → {new_w}×{new_h}")
     return image.resize((new_w, new_h), Image.LANCZOS)
 
 
-def transform_image(image: Image.Image, resolution: int = 512) -> torch.Tensor:
-    image = image.resize((resolution, resolution))
-    arr = np.array(image).astype(np.float32) / 255.0
+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)
-    arr = (arr - mean) / std
 
+    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
 
 
-def run_inference(image: Image.Image, resolution: int = 512) -> Image.Image:
+def run_inference(image: Image.Image) -> Image.Image:
     orig_size = image.size
-
-    input_tensor = transform_image(image, resolution)
+    input_tensor = transform_image(image)
 
     with inference_lock:
         with torch.no_grad():
@@ -102,24 +108,21 @@ def run_inference(image: Image.Image, resolution: int = 512) -> Image.Image:
     image.putalpha(mask)
     return image
 
-
 # ---------------------------------------------------------
 # FastAPI app
 # ---------------------------------------------------------
 app = FastAPI(title="Background Remover API")
 
-
 # ---------------------------------------------------------
-# remove-background endpoint
+# POST endpoint only (no GET processing)
 # ---------------------------------------------------------
 @app.post("/remove-background")
 async def remove_background(
     file: UploadFile = File(None),
-    image_url: str = Form(None),
-    resolution: int = Form(512)
+    image_url: str = Form(None)
 ):
     try:
-        # Load input
+        # load image
         if file:
             raw = await file.read()
             image = Image.open(BytesIO(raw)).convert("RGB")
@@ -128,12 +131,13 @@ async def remove_background(
         else:
             raise HTTPException(status_code=400, detail="Provide file or image_url.")
 
-        # Automatically compress very large images
+        # auto shrink large inputs
         image = auto_downscale(image)
 
-        # Process image
-        result = run_inference(image, resolution)
+        # remove background
+        result = run_inference(image)
 
+        # return PNG
         buf = BytesIO()
         result.save(buf, format="PNG", optimize=True)
         buf.seek(0)
@@ -145,78 +149,74 @@ async def remove_background(
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")
 
-
 # ---------------------------------------------------------
-# Web test UI
+# UI for POST method testing only
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 async def ui():
     return """
     <html>
     <head>
-        <title>Background Remover</title>
+        <title>Background Remover – Test Tool</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>Background Remover API</h2>
+            <h2>POST Method Test Panel</h2>
+            <p>This UI only sends POST requests to <code>/remove-background</code>.</p>
 
-            <form id='f1' class='mb-4' enctype='multipart/form-data'>
-                <input class='form-control mb-2' type='file' id='fileInput' name='file'>
-                <input class='form-control mb-2' type='number' id='resInput' name='resolution' value='512'>
-                <button class='btn btn-primary'>Upload</button>
+            <h5>Test with File Upload:</h5>
+            <form id='uploadForm' enctype='multipart/form-data'>
+                <input class='form-control mb-2' type='file' id='fileInput'>
+                <button class='btn btn-primary'>Send POST</button>
             </form>
 
-            <div class='mb-3'>OR</div>
+            <div class='my-4'>OR</div>
 
-            <form id='f2'>
-                <input class='form-control mb-2' id='urlInput' placeholder='Image URL'>
-                <input class='form-control mb-2' id='urlResInput' type='number' value='512'>
-                <button class='btn btn-success'>Use URL</button>
+            <h5>Test with Image URL:</h5>
+            <form id='urlForm'>
+                <input class='form-control mb-2' id='urlInput' placeholder='https://example.com/image.jpg'>
+                <button class='btn btn-success'>Send POST</button>
             </form>
 
-            <h5 class='mt-4'>Result:</h5>
-            <img id='out' style='max-width:100%;border-radius:10px;'/>
+            <h4 class='mt-4'>Output:</h4>
+            <img id='outputImg' style='max-width:90%;border-radius:10px;'/>
         </div>
 
         <script>
-            const out = document.getElementById("out");
+            const outputImg = document.getElementById("outputImg");
 
-            document.getElementById("f1").addEventListener("submit", async e => {
+            document.getElementById("uploadForm").addEventListener("submit", async e => {
                 e.preventDefault();
                 const file = document.getElementById("fileInput").files[0];
-                if (!file) return alert("Select an image");
-                const res = document.getElementById("resInput").value;
+                if (!file) return alert("Select a file first");
 
                 const fd = new FormData();
                 fd.append("file", file);
-                fd.append("resolution", res);
 
-                const r = await fetch("/remove-background", { method:"POST", body:fd });
-                out.src = URL.createObjectURL(await r.blob());
+                const r = await fetch("/remove-background", {method:"POST", body:fd});
+                outputImg.src = URL.createObjectURL(await r.blob());
             });
 
-            document.getElementById("f2").addEventListener("submit", async e => {
+            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 res = document.getElementById("urlResInput").value;
 
                 const fd = new FormData();
                 fd.append("image_url", url);
-                fd.append("resolution", res);
 
-                const r = await fetch("/remove-background", { method:"POST", body:fd });
-                out.src = URL.createObjectURL(await r.blob());
+                const r = await fetch("/remove-background", {method:"POST", body:fd});
+                outputImg.src = URL.createObjectURL(await r.blob());
             });
         </script>
     </body>
     </html>
     """
 
-
 # ---------------------------------------------------------
-# Start server
+# Run app
 # ---------------------------------------------------------
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)