Create app.py

app.py

@@ -0,0 +1,59 @@
+import io
+import os
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import StreamingResponse
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+import onnxruntime as ort
+
+# Settings
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+ONNX_PATH = os.path.join(os.path.dirname(__file__), "birefnet.onnx")
+
+# Preprocessing transform
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+# Load ONNX model
+if not os.path.exists(ONNX_PATH):
+    raise FileNotFoundError(f"ONNX model not found at {ONNX_PATH}")
+
+providers = ["CUDAExecutionProvider"] if DEVICE == "cuda" else ["CPUExecutionProvider"]
+onnx_session = ort.InferenceSession(ONNX_PATH, providers=providers)
+print(f"ONNX model loaded with providers: {providers}")
+
+# Helper functions
+def run_model_onnx(input_tensor: torch.Tensor) -> torch.Tensor:
+    ort_inputs = {onnx_session.get_inputs()[0].name: input_tensor.cpu().numpy()}
+    ort_outs = onnx_session.run(None, ort_inputs)
+    preds = torch.from_numpy(ort_outs[-1]).sigmoid()
+    return preds
+
+def process_image(image: Image.Image) -> Image.Image:
+    original_size = image.size
+    input_tensor = transform_image(image).unsqueeze(0)  # (1,C,H,W)
+    preds = run_model_onnx(input_tensor)
+    pred = preds[0]
+    if pred.dim() == 3:
+        pred = pred[0].squeeze(0)
+    mask = transforms.ToPILImage()(pred.clamp(0, 1))
+    mask = mask.resize(original_size, resample=Image.BILINEAR)
+    image_rgba = image.convert("RGBA")
+    image_rgba.putalpha(mask)
+    return image_rgba
+
+# FastAPI app
+app = FastAPI(title="Background Removal API")
+
+@app.post("/remove-background")
+async def remove_background(file: UploadFile = File(...)):
+    image = Image.open(file.file).convert("RGB")
+    result_image = process_image(image)
+    buf = io.BytesIO()
+    result_image.save(buf, format="PNG")
+    buf.seek(0)
+    return StreamingResponse(buf, media_type="image/png")