Background Remover API Test

import os from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Query 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/HEIF Support # ------------------------- try: import pillow_heif pillow_heif.register_heif_opener() print("✅ HEIC/HEIF format supported.") except ImportError: print("⚠️ Install pillow-heif for HEIC support: pip install pillow-heif") # ------------------------- # Model Setup # ------------------------- MODEL_DIR = "models/BiRefNet" 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 BiRefNet model...") birefnet = AutoModelForImageSegmentation.from_pretrained( "ZhengPeng7/BiRefNet", cache_dir=MODEL_DIR, trust_remote_code=True, revision="main" ) birefnet.to(device, dtype=dtype).eval() print("Model loaded successfully.") # ------------------------- # FastAPI App # ------------------------- app = FastAPI(title="Background Remover API") # ------------------------- # Utility Functions # ------------------------- 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, resolution: int = 512) -> torch.Tensor: 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 = np.transpose(arr, (2, 0, 1)) # HWC -> CHW tensor = torch.from_numpy(arr).unsqueeze(0).to(dtype).to(device) return tensor def process_image(image: Image.Image, resolution: int = 512) -> Image.Image: orig_size = image.size input_tensor = transform_image(image, resolution) with torch.no_grad(): preds = birefnet(input_tensor)[-1].sigmoid().cpu() 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 # ------------------------- # /remove-background Endpoint # ------------------------- @app.post("/remove-background") async def remove_background( file: UploadFile = File(None), image_url: str = Form(None), resolution: int = Form(512) ): """ Remove background from an image. Accepts a file upload or image URL. Optional resolution (default 512) for faster inference. Returns 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, resolution) 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)) # ------------------------- # Developer Test Page (Bootstrap) # ------------------------- @app.get("/", response_class=HTMLResponse) async def index(): html = """ Background Remover API Test

Background Remover API Tester

Result:

""" return HTMLResponse(html) # ------------------------- # Run App # ------------------------- if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=7860)