"""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 hmac 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 Depends, FastAPI, File, Form, Header, 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 # Optional app-level auth. When BGFILTER_AUTH_TOKEN is set, protected endpoints # require the X-Genarrative-Image-Token header to equal it (constant-time compare) -- # the same header the production nginx uses, so callers only swap the URL + token. # When unset the check is a no-op (open), matching pre-auth behaviour and leaving # auth to an upstream proxy (e.g. nginx in production). AUTH_TOKEN = os.environ.get("BGFILTER_AUTH_TOKEN", "").strip() async def require_auth(x_genarrative_image_token: str | None = Header(default=None)): if not AUTH_TOKEN: return if not x_genarrative_image_token or not hmac.compare_digest( x_genarrative_image_token, AUTH_TOKEN ): raise HTTPException(status_code=401, detail="missing or invalid auth token") _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) # The cross-model veto is default behaviour (config default: on). With # reuse-as-seg it replaces the primary seg forward rather than adding to it, # so the earlier CPU-cost concern no longer applies. BGFILTER_CROSS_CHECK # can still force it either way; unset = honour the config. cc_env = os.environ.get("BGFILTER_CROSS_CHECK") if cc_env is not None and cc_env.strip() != "": cross_check = cc_env.strip().lower() in ("1", "true", "yes", "on") settings = override_settings(settings, cross_check=cross_check) 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 cross_check=%s preload=%s max_pixels=%d auth=%s", settings.model.device, settings.segmentation.backend, settings.cross_check.enabled, preload, MAX_IMAGE_PIXELS, "enabled" if AUTH_TOKEN else "DISABLED(open)", ) if not AUTH_TOKEN: logger.warning( "BGFILTER_AUTH_TOKEN not set: /remove-background is OPEN (no app-level auth)" ) _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, "crossCheck": _manager.cross_check_enabled, } 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 _normalize_cross_check(value: str | None) -> bool | None: """None/empty -> use the service default; otherwise a boolean string.""" if value is None: return None v = value.strip().lower() if v == "": return None if v in ("1", "true", "yes", "on"): return True if v in ("0", "false", "no", "off"): return False raise HTTPException(status_code=400, detail="invalid cross_check") 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), cross_check: str | None = Form(None), _auth: None = Depends(require_auth), ): 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 normalized_cc = _normalize_cross_check(cross_check) effective_cc = normalized_cc if normalized_cc is not None else _manager.cross_check_enabled 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, cross_check=normalized_cc, ), ) 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", "X-BGFilter-Cross-Check": "on" if effective_cc else "off", } return Response(content=png, media_type="image/png", headers=headers)