BGfilter/app.py

"""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)