Spaces:

videopix
/

background-remover-api

Running

App Files Files

videopix commited on Oct 16, 2025

Commit

b88e1b6

verified ·

1 Parent(s): f510e5e

Update app.py

Browse files

Files changed (1) hide show

app.py +33 -20

app.py CHANGED Viewed

@@ -1,13 +1,14 @@
 import os
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import StreamingResponse, HTMLResponse
-from PIL import Image
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
 from io import BytesIO
 import requests
 import uvicorn
 # -------------------------
 # Optional HEIC Support
@@ -33,13 +34,17 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
     trust_remote_code=True,
     revision="main"
 )
-birefnet.to(device).eval()
 print("Model loaded successfully.")
 # -------------------------
 # FastAPI App
 # -------------------------
 app = FastAPI(title="Background Remover API")
 # -------------------------
 # Utility Functions
@@ -48,18 +53,18 @@ def load_image_from_url(url: str) -> Image.Image:
     try:
         response = requests.get(url, timeout=10)
         response.raise_for_status()
-        return Image.open(BytesIO(response.content)).convert("RGB")
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Error loading image from URL: {str(e)}")
-def transform_image(image: Image.Image) -> torch.Tensor:
-    image = image.resize((1024, 1024))
     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 = np.transpose(arr, (2, 0, 1))  # HWC -> CHW
-    tensor = torch.from_numpy(arr).unsqueeze(0).to(torch.float32).to(device)
     return tensor
 def process_image(image: Image.Image) -> Image.Image:
@@ -80,20 +85,27 @@ def process_image(image: Image.Image) -> Image.Image:
 async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
     """
     Remove background from an image.
-    Accepts either a file upload or an image URL.
-    Returns a PNG with transparent background.
     """
     try:
         if file:
-            image = Image.open(BytesIO(await file.read())).convert("RGB")
         elif image_url:
             image = load_image_from_url(image_url)
         else:
-            raise HTTPException(status_code=400, detail="Provide either 'file' or 'image_url'.")
-        result = process_image(image)
-        buf = BytesIO()
-        result.save(buf, format="PNG")
         buf.seek(0)
         return StreamingResponse(buf, media_type="image/png")
     except Exception as e:
@@ -121,35 +133,32 @@ async def index():
     <body>
         <div class="container text-center">
             <h2 class="mb-4">Background Remover API Tester</h2>
             <form id="uploadForm" class="mb-4" enctype="multipart/form-data">
                 <div class="mb-3">
-                    <label for="fileInput" class="form-label">Upload Image (any format, e.g. JPG, PNG, HEIC):</label>
                     <input class="form-control" type="file" id="fileInput" name="file" accept="image/*">
                 </div>
                 <button class="btn btn-primary" type="submit">Remove Background</button>
             </form>
             <div class="mb-4">OR</div>
             <form id="urlForm" class="mb-4">
                 <div class="mb-3">
-                    <label for="urlInput" class="form-label">Enter Image URL:</label>
                     <input class="form-control" type="text" id="urlInput" placeholder="https://example.com/image.jpg">
                 </div>
                 <button class="btn btn-success" type="submit">Remove Background</button>
             </form>
             <div id="resultContainer" class="mt-4">
                 <h5>Result:</h5>
                 <img id="resultImg" src="" alt="">
             </div>
         </div>
         <script>
             const uploadForm = document.getElementById("uploadForm");
             const urlForm = document.getElementById("urlForm");
             const resultImg = document.getElementById("resultImg");
             uploadForm.addEventListener("submit", async e => {
                 e.preventDefault();
@@ -160,6 +169,8 @@ async def index():
                 const res = await fetch("/remove-background", { method: "POST", body: formData });
                 const blob = await res.blob();
                 resultImg.src = URL.createObjectURL(blob);
             });
             urlForm.addEventListener("submit", async e => {
@@ -171,6 +182,8 @@ async def index():
                 const res = await fetch("/remove-background", { method: "POST", body: formData });
                 const blob = await res.blob();
                 resultImg.src = URL.createObjectURL(blob);
             });
         </script>
     </body>

 import os
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import StreamingResponse, HTMLResponse
+from PIL import Image, ImageSequence
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
 from io import BytesIO
 import requests
 import uvicorn
+from concurrent.futures import ThreadPoolExecutor
 # -------------------------
 # Optional HEIC Support
     trust_remote_code=True,
     revision="main"
 )
+birefnet.to(device)
+if device == "cuda":
+    birefnet = birefnet.half()  # FP16 for faster GPU inference
+birefnet.eval()
 print("Model loaded successfully.")
 # -------------------------
 # FastAPI App
 # -------------------------
 app = FastAPI(title="Background Remover API")
+executor = ThreadPoolExecutor(max_workers=4)
 # -------------------------
 # Utility Functions
     try:
         response = requests.get(url, timeout=10)
         response.raise_for_status()
+        return Image.open(BytesIO(response.content))
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Error loading image from URL: {str(e)}")
+def transform_image(image: Image.Image, target_size=(512, 512)) -> 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 = np.transpose(arr, (2, 0, 1))  # HWC -> CHW
+    tensor = torch.from_numpy(arr).unsqueeze(0).to(torch.float16 if device=="cuda" else torch.float32).to(device)
     return tensor
 def process_image(image: Image.Image) -> Image.Image:
 async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
     """
     Remove background from an image.
+    Accepts file upload or image URL.
+    Returns PNG with transparent background.
     """
     try:
         if file:
+            image = Image.open(BytesIO(await file.read()))
         elif image_url:
             image = load_image_from_url(image_url)
         else:
+            raise HTTPException(status_code=400, detail="Provide 'file' or 'image_url'.")
+        # Handle multi-frame images (GIF, PDF)
+        if getattr(image, "is_animated", False):
+            frames = [process_image(frame.convert("RGBA")) for frame in ImageSequence.Iterator(image)]
+            buf = BytesIO()
+            frames[0].save(buf, format="PNG", save_all=True, append_images=frames[1:])
+        else:
+            result = process_image(image.convert("RGBA"))
+            buf = BytesIO()
+            result.save(buf, format="PNG")
         buf.seek(0)
         return StreamingResponse(buf, media_type="image/png")
     except Exception as e:
     <body>
         <div class="container text-center">
             <h2 class="mb-4">Background Remover API Tester</h2>
             <form id="uploadForm" class="mb-4" enctype="multipart/form-data">
                 <div class="mb-3">
+                    <label class="form-label">Upload Image (any format):</label>
                     <input class="form-control" type="file" id="fileInput" name="file" accept="image/*">
                 </div>
                 <button class="btn btn-primary" type="submit">Remove Background</button>
             </form>
             <div class="mb-4">OR</div>
             <form id="urlForm" class="mb-4">
                 <div class="mb-3">
+                    <label class="form-label">Enter Image URL:</label>
                     <input class="form-control" type="text" id="urlInput" placeholder="https://example.com/image.jpg">
                 </div>
                 <button class="btn btn-success" type="submit">Remove Background</button>
             </form>
             <div id="resultContainer" class="mt-4">
                 <h5>Result:</h5>
                 <img id="resultImg" src="" alt="">
+                <a id="downloadLink" class="btn btn-info mt-2" download="result.png" style="display:none;">Download PNG</a>
             </div>
         </div>
         <script>
             const uploadForm = document.getElementById("uploadForm");
             const urlForm = document.getElementById("urlForm");
             const resultImg = document.getElementById("resultImg");
+            const downloadLink = document.getElementById("downloadLink");
             uploadForm.addEventListener("submit", async e => {
                 e.preventDefault();
                 const res = await fetch("/remove-background", { method: "POST", body: formData });
                 const blob = await res.blob();
                 resultImg.src = URL.createObjectURL(blob);
+                downloadLink.href = resultImg.src;
+                downloadLink.style.display = "inline-block";
             });
             urlForm.addEventListener("submit", async e => {
                 const res = await fetch("/remove-background", { method: "POST", body: formData });
                 const blob = await res.blob();
                 resultImg.src = URL.createObjectURL(blob);
+                downloadLink.href = resultImg.src;
+                downloadLink.style.display = "inline-block";
             });
         </script>
     </body>