Add FastAPI HTTP service for background removal

Wrap the BGfilter matting pipeline in a long-running FastAPI service (POST /remove-background, GET /healthz) exposing only screen_color and seg_model, matching the birefnet-service request contract. - bgfilter/pipeline.py: add in-memory run_rgb entry returning a MattingResult; allow injecting a shared runner/segmenter. CLI file path (run_image) behaviour is unchanged. - bgfilter/service.py: PipelineManager caches one shared ViTMatte runner and per-backend segmenters, so screen_color/seg_model never reload a model. - app.py: thin FastAPI layer -- param validation, image decode/encode, error codes (400/413/500), timing headers, global asyncio lock. - requirements.txt: fastapi, uvicorn[standard], python-multipart. - README + design doc: HTTP service usage and HF-mirror/proxy/Xet download notes. Verified end-to-end (base ViTMatte + BiRefNet + anime-seg): 11/11 acceptance checks pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-02 12:08:25 +00:00
parent cf0ceec496
commit d24c6ed608
6 changed files with 892 additions and 33 deletions
@@ -0,0 +1,196 @@
+"""FastAPI entry for the BGfilter background-removal service.
+
+Thin HTTP layer: parameter validation, image decode/encode, error codes, timing,
+and a global lock. All model handling lives in ``bgfilter.service.PipelineManager``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import io
+import logging
+import os
+import re
+import time
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+import numpy as np
+from fastapi import FastAPI, File, Form, HTTPException, Response, UploadFile
+from PIL import Image, ImageOps, UnidentifiedImageError
+
+from bgfilter.config import load_settings, override_settings
+from bgfilter.service import SEG_MODELS, PipelineManager
+
+logging.basicConfig(
+    level=os.environ.get("BGFILTER_LOG_LEVEL", "INFO"),
+    format="%(asctime)s %(levelname)s %(name)s %(message)s",
+)
+logger = logging.getLogger("bgfilter.app")
+
+VERSION = "0.1.0"
+_HEX_RE = re.compile(r"^#[0-9a-fA-F]{6}$")
+
+
+def _int_env(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if not raw:
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        logger.warning("ignoring non-integer %s=%r; using %d", name, raw, default)
+        return default
+
+
+MAX_IMAGE_PIXELS = _int_env("BGFILTER_MAX_IMAGE_PIXELS", 4_194_304)
+# We enforce our own pixel budget from the image header before decoding, so disable
+# Pillow's separate decompression-bomb ceiling (which would otherwise raise its own error).
+Image.MAX_IMAGE_PIXELS = None
+
+_manager: PipelineManager | None = None
+_lock = asyncio.Lock()
+
+
+def _build_settings():
+    config_path = os.environ.get("BGFILTER_CONFIG")
+    if config_path is None:
+        default_cfg = Path("configs/default.yaml")
+        config_path = str(default_cfg) if default_cfg.exists() else None
+    settings = load_settings(config_path)
+    device = os.environ.get("BGFILTER_DEVICE")
+    if device:
+        # Covers both settings.model.device and settings.segmentation.device.
+        settings = override_settings(settings, device=device)
+    return settings
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global _manager
+    settings = _build_settings()
+    preload = os.environ.get("BGFILTER_PRELOAD", "1").strip().lower() not in ("0", "false", "no", "")
+    logger.info(
+        "starting bgfilter service: device=%s default_seg=%s preload=%s max_pixels=%d",
+        settings.model.device,
+        settings.segmentation.backend,
+        preload,
+        MAX_IMAGE_PIXELS,
+    )
+    _manager = PipelineManager(settings, preload=preload)
+    yield
+
+
+app = FastAPI(title="BGfilter", version=VERSION, lifespan=lifespan)
+
+
+@app.get("/healthz")
+def healthz():
+    if _manager is None:
+        raise HTTPException(status_code=503, detail="service not ready")
+    return {
+        "ok": True,
+        "service": "bgfilter",
+        "version": VERSION,
+        "defaultSegModel": _manager.default_seg_model,
+        "device": _manager.device,
+    }
+
+
+def _normalize_screen_color(value: str | None) -> str | None:
+    """None/empty/"null"/"auto" -> auto-detect; otherwise require #RRGGBB."""
+    if value is None:
+        return None
+    v = value.strip()
+    if v == "" or v.lower() in ("null", "auto"):
+        return None
+    if not _HEX_RE.match(v):
+        raise HTTPException(status_code=400, detail="invalid screen_color")
+    return v.upper()
+
+
+def _normalize_seg_model(value: str | None) -> str | None:
+    """None/empty -> use the service default; otherwise must be a known backend."""
+    if value is None:
+        return None
+    v = value.strip().lower()
+    if v == "":
+        return None
+    if v not in SEG_MODELS:
+        raise HTTPException(status_code=400, detail="invalid seg_model")
+    return v
+
+
+def _decode_image(data: bytes) -> np.ndarray:
+    if not data:
+        raise HTTPException(status_code=400, detail="empty image")
+    try:
+        image = Image.open(io.BytesIO(data))
+    except (UnidentifiedImageError, OSError, ValueError):
+        raise HTTPException(status_code=400, detail="invalid image")
+    width, height = image.size
+    if width <= 0 or height <= 0:
+        raise HTTPException(status_code=400, detail="invalid image")
+    if width * height > MAX_IMAGE_PIXELS:
+        raise HTTPException(status_code=413, detail="image too large")
+    try:
+        image = ImageOps.exif_transpose(image).convert("RGB")
+        image.load()  # force full decode; truncated files raise here
+    except (OSError, ValueError):
+        raise HTTPException(status_code=400, detail="invalid image")
+    return np.asarray(image, dtype=np.uint8)
+
+
+def _encode_png(rgb: np.ndarray, alpha: np.ndarray) -> bytes:
+    alpha_u8 = np.clip(alpha * 255.0, 0, 255).astype(np.uint8)
+    rgba = np.dstack([np.clip(rgb, 0, 255).astype(np.uint8), alpha_u8])
+    buffer = io.BytesIO()
+    Image.fromarray(rgba, mode="RGBA").save(buffer, format="PNG")
+    return buffer.getvalue()
+
+
+@app.post("/remove-background")
+async def remove_background(
+    file: UploadFile = File(...),
+    screen_color: str | None = Form(None),
+    seg_model: str | None = Form(None),
+):
+    if _manager is None:
+        raise HTTPException(status_code=503, detail="service not ready")
+
+    normalized_color = _normalize_screen_color(screen_color)
+    normalized_seg = _normalize_seg_model(seg_model) or _manager.default_seg_model
+
+    rgb = _decode_image(await file.read())
+
+    start = time.perf_counter()
+    async with _lock:  # one inference at a time per process
+        loop = asyncio.get_running_loop()
+        try:
+            result = await loop.run_in_executor(
+                None,
+                lambda: _manager.process(
+                    rgb, screen_color=normalized_color, seg_model=normalized_seg
+                ),
+            )
+        except RuntimeError as exc:
+            # No clean flat background and no screen_color given -> client-actionable.
+            if "auto-detect" in str(exc).lower():
+                raise HTTPException(
+                    status_code=400,
+                    detail="could not auto-detect background colour; pass screen_color",
+                )
+            logger.exception("inference failed")
+            raise HTTPException(status_code=500, detail="inference failed")
+        except Exception:
+            logger.exception("inference failed")
+            raise HTTPException(status_code=500, detail="inference failed")
+    elapsed_ms = int((time.perf_counter() - start) * 1000)
+
+    png = _encode_png(result.rgb, result.alpha)
+    headers = {
+        "X-BGFilter-Elapsed-Ms": str(elapsed_ms),
+        "X-BGFilter-Seg-Model": normalized_seg,
+        "X-BGFilter-Screen-Color": normalized_color if normalized_color else "auto",
+    }
+    return Response(content=png, media_type="image/png", headers=headers)