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 contextvars
import datetime
import hmac
import io
import json
import logging
import os
import re
import time
import traceback
import uuid
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from logging.handlers import RotatingFileHandler
from pathlib import Path
from urllib.parse import urlparse

import httpx
import numpy as np
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Request, Response, UploadFile
from PIL import Image, ImageOps, UnidentifiedImageError

from bgfilter.config import load_settings, override_settings
from bgfilter.service import SEG_MODELS, PipelineManager

# Per-request correlation id. A middleware sets it; a logging filter injects it into
# every record so all logs for one request share the id. contextvars (not thread-
# locals) so a pooled worker thread never inherits the previous request's id.
_request_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("request_id", default="-")
_RID_RE = re.compile(r"^[A-Za-z0-9._-]{1,64}$")


def _make_request_id(incoming: str | None) -> str:
    """Reuse a sanitized inbound X-Request-ID (e.g. nginx $request_id), else generate."""
    if incoming and _RID_RE.match(incoming):
        return incoming
    return uuid.uuid4().hex[:16]


class _RequestIdFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        record.request_id = _request_id_var.get()
        return True


class _JsonFormatter(logging.Formatter):
    """One JSON object per record; per-request access fields ride in `extra={"fields"}`."""

    def format(self, record: logging.LogRecord) -> str:
        out = {
            "ts": datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z"),
            "level": record.levelname,
            "logger": record.name,
            "request_id": getattr(record, "request_id", "-"),
            "msg": record.getMessage(),
        }
        fields = getattr(record, "fields", None)
        if isinstance(fields, dict):
            out.update(fields)
        if record.exc_info:
            out["exc"] = self.formatException(record.exc_info)
        return json.dumps(out, ensure_ascii=False)


# Two log streams with distinct jobs (never mixed into one file):
#
# * Event stream -- prose for humans, low volume: startup summary, warnings, error
#   tracebacks. Goes to stderr (text format) where supervisord/journald captures it.
#   Per-request records must NOT be logged here: bulk structured content does not
#   belong in the prose file.
# * Access stream -- structured JSONL, one record per request (plus one "startup"
#   meta record): everything with analysis value. Written to the file named by
#   BGFILTER_ACCESS_LOG (rotated by the app itself: supervisord can't rotate a file
#   it doesn't own); when unset (dev), the JSON lines fall back to stderr.
#
# Both streams stamp the same request_id, so a traceback in the event stream joins
# back to its access record (which also carries a truncated `exc` copy).
_event_handler = logging.StreamHandler()
_event_handler.addFilter(_RequestIdFilter())
_event_handler.setFormatter(
    logging.Formatter("%(asctime)s %(levelname)s %(name)s [%(request_id)s] %(message)s")
)
logging.basicConfig(level=os.environ.get("BGFILTER_LOG_LEVEL", "INFO"), handlers=[_event_handler])
# httpx logs one INFO line per outbound request (full presigned image_url included);
# that is per-request noise on the prose stream -- the access JSONL already records
# the URL. Keep only its warnings/errors.
logging.getLogger("httpx").setLevel(logging.WARNING)
logger = logging.getLogger("bgfilter.app")

ACCESS_LOG_PATH = os.environ.get("BGFILTER_ACCESS_LOG", "").strip()
access_logger = logging.getLogger("bgfilter.access")
access_logger.propagate = False  # never leak per-request records into the prose stream
if ACCESS_LOG_PATH:
    Path(ACCESS_LOG_PATH).parent.mkdir(parents=True, exist_ok=True)
    _access_handler: logging.Handler = RotatingFileHandler(
        ACCESS_LOG_PATH, maxBytes=50 * 1024 * 1024, backupCount=5, encoding="utf-8"
    )
else:
    _access_handler = logging.StreamHandler()
_access_handler.addFilter(_RequestIdFilter())
_access_handler.setFormatter(_JsonFormatter())
access_logger.addHandler(_access_handler)


def _emit_access(event: str, fields: dict) -> None:
    """One JSONL record on the access stream (`msg` = record type)."""
    access_logger.info(event, extra={"fields": fields})

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)
# /remove-background accepts either an uploaded `file` or an `image_url` the service
# fetches itself. These bound the fetch so a huge or slow URL cannot hang a worker or
# blow RAM. NOTE: no SSRF hardening here (no host/IP allowlist) -- the endpoint is
# guarded by token auth and assumes trusted callers; do NOT expose it to untrusted
# clients without adding private-address / redirect filtering to _fetch_image.
FETCH_MAX_BYTES = _int_env("BGFILTER_FETCH_MAX_BYTES", 25 * 1024 * 1024)
FETCH_TIMEOUT_S = float(os.environ.get("BGFILTER_FETCH_TIMEOUT", "15"))
# Outbound fetch concurrency = 2 x cpu_workers, hard-capped here so a burst can never
# exhaust file descriptors / ephemeral ports (each fetch = one socket).
FETCH_CONCURRENCY_CAP = 64
# 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

# Request pipelining. Each request's process() runs on a small thread pool so the
# CPU stages (chroma / trimap / alpha_post / foreground) of concurrent requests
# overlap, while a GPU mutex inside the pipeline keeps the model forwards mutually
# exclusive on the single device -- i.e. when one request steps off the GPU to do
# CPU work, another steps on, keeping the GPU fed instead of idle.
# The pool size defaults to config (server.cpu_workers, default 1 == strictly serial)
# and is overridable by BGFILTER_CPU_WORKERS. Both are resolved in lifespan once the
# config is loaded; a same-sized semaphore bounds in-flight requests (hence
# decoded-image RAM) so a burst cannot pile unbounded decoded images before the pool
# drains.
_executor: ThreadPoolExecutor | None = None
_sem: asyncio.Semaphore | None = None
# Outbound image_url fetches: a FIFO admission bound (restores the throttle a direct
# upload has implicitly) and a shared client for connection reuse. Both built in
# lifespan once the pool size is known. See _fetch_image / §image_url mode.
_fetch_sem: asyncio.Semaphore | None = None
_http_client: httpx.AsyncClient | None = None


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, _executor, _sem, _fetch_sem, _http_client
    settings = _build_settings()
    preload = os.environ.get("BGFILTER_PRELOAD", "1").strip().lower() not in ("0", "false", "no", "")
    # Pool size: config default (server.cpu_workers), env override wins.
    cpu_workers = max(1, _int_env("BGFILTER_CPU_WORKERS", settings.server.cpu_workers))
    _executor = ThreadPoolExecutor(max_workers=cpu_workers, thread_name_prefix="bgf")
    _sem = asyncio.Semaphore(cpu_workers)
    # Outbound-fetch bound: 2x the pool (a small prefetch buffer), hard-capped. The
    # shared client's max_connections matches so httpx never opens more sockets than
    # the semaphore admits, and reuses connections across fetches.
    fetch_conc = min(2 * cpu_workers, FETCH_CONCURRENCY_CAP)
    _fetch_sem = asyncio.Semaphore(fetch_conc)
    _http_client = httpx.AsyncClient(
        follow_redirects=True,
        max_redirects=3,
        timeout=httpx.Timeout(connect=5.0, read=FETCH_TIMEOUT_S, write=10.0, pool=None),
        limits=httpx.Limits(max_connections=fetch_conc),
    )
    gpu_concurrency = max(1, _int_env("BGFILTER_GPU_CONCURRENCY", settings.server.gpu_concurrency))
    logger.info(
        "starting bgfilter service: device=%s default_seg=%s cross_check=%s preload=%s "
        "cpu_workers=%d gpu_concurrency=%d max_pixels=%d auth=%s",
        settings.model.device,
        settings.segmentation.backend,
        settings.cross_check.enabled,
        preload,
        cpu_workers,
        gpu_concurrency,
        MAX_IMAGE_PIXELS,
        "enabled" if AUTH_TOKEN else "DISABLED(open)",
    )
    # Same facts as a structured record on the access stream, so per-instance config
    # is queryable next to the requests it served (no prose parsing).
    _emit_access(
        "startup",
        {
            "version": VERSION,
            "device": settings.model.device,
            "default_seg": settings.segmentation.backend,
            "cross_check": settings.cross_check.enabled,
            "reuse_as_seg": settings.cross_check.reuse_as_seg,
            "preload": preload,
            "cpu_workers": cpu_workers,
            "gpu_concurrency": gpu_concurrency,
            "max_pixels": MAX_IMAGE_PIXELS,
            "auth": "enabled" if AUTH_TOKEN else "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)
    try:
        yield
    finally:
        await _http_client.aclose()


app = FastAPI(title="BGfilter", version=VERSION, lifespan=lifespan)


@app.middleware("http")
async def _access_log_mw(request: Request, call_next):
    """Bind a request id, echo it back, and emit one structured access line per request.

    The handler accumulates its fields on ``request.state.fields``; here we add the
    id, client, status and total time -- so even a request that errors out (the
    handler raised before finishing) still produces a log line.
    """
    rid = _make_request_id(request.headers.get("x-request-id"))
    _request_id_var.set(rid)
    request.state.fields = {}
    t0 = time.perf_counter()
    status = 500
    try:
        response = await call_next(request)
        status = response.status_code
        response.headers["X-Request-ID"] = rid
        return response
    except Exception:
        # Unhandled escape (HTTPExceptions were already turned into responses by the
        # inner exception middleware). setdefault: a more precise traceback recorded
        # at the raise site wins over this generic one.
        request.state.fields.setdefault("exc", traceback.format_exc()[-2000:])
        raise
    finally:
        if request.url.path != "/healthz":
            fields = {"request_id": rid, "client": request.client.host if request.client else "-"}
            fields.update(getattr(request.state, "fields", {}) or {})
            fields["status"] = status
            fields["t_total_ms"] = round((time.perf_counter() - t0) * 1000)
            _emit_access("request", fields)


@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 _normalize_background_mode(value: str | None) -> str | None:
    """None/empty -> service default (flat); else "flat" or "complex".

    "flat" (default) is the colour-keyed pipeline; "complex" is non-flat/scene
    matting with no colour key (segmentation drives the trimap). We deliberately
    do not surface the internal "chroma" wording here.
    """
    if value is None:
        return None
    v = value.strip().lower()
    if v == "":
        return None
    if v in ("flat", "complex"):
        return v
    raise HTTPException(status_code=400, detail="invalid background_mode")


async def _fetch_image(url: str) -> bytes:
    """Download image bytes from a URL, bounded by size and time.

    Robustness only -- no SSRF filtering (see FETCH_* note above): callers are
    trusted via the auth token. Streams the body so an over-large response is
    rejected without being fully buffered. Uses the shared client (connection reuse);
    the caller holds _fetch_sem so outbound concurrency stays bounded.
    """
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https") or not parsed.hostname:
        raise HTTPException(status_code=400, detail="image_url must be an http(s) URL")
    try:
        async with _http_client.stream("GET", url) as resp:
            if resp.status_code != 200:
                raise HTTPException(
                    status_code=400,
                    detail=f"image_url fetch failed (status {resp.status_code})",
                )
            buf = bytearray()
            async for chunk in resp.aiter_bytes():
                buf += chunk
                if len(buf) > FETCH_MAX_BYTES:
                    raise HTTPException(status_code=413, detail="image_url body too large")
            return bytes(buf)
    except httpx.HTTPError:
        raise HTTPException(status_code=400, detail="image_url fetch failed")


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(
    request: Request,
    file: UploadFile | None = File(None),
    image_url: str | None = Form(None),
    screen_color: str | None = Form(None),
    seg_model: str | None = Form(None),
    cross_check: str | None = Form(None),
    background_mode: str | None = Form(None),
    _auth: None = Depends(require_auth),
):
    if _manager is None:
        raise HTTPException(status_code=503, detail="service not ready")

    # Access-log fields accumulate here; the middleware emits the line (with status /
    # total time) even if we raise below. Record the error detail on the way out.
    log = request.state.fields
    try:
        # Exactly one image source. (Auth already ran, so only authenticated callers
        # can trigger an outbound fetch.)
        url = image_url.strip() if image_url else None
        if (file is None) == (not url):
            raise HTTPException(
                status_code=400, detail="provide exactly one of 'file' or 'image_url'"
            )

        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)
        # None/"flat" -> flat (chroma on); "complex" -> non-flat (chroma off).
        normalized_bg = _normalize_background_mode(background_mode)
        chroma_enabled = normalized_bg != "complex"
        effective_bg = "complex" if not chroma_enabled else "flat"
        # Complex mode defaults the cross-check veto OFF (it costs the HR-matting
        # forward and only clears residue near strands); an explicit cross_check=on
        # still enables its hue-free gate.
        if not chroma_enabled and normalized_cc is None:
            normalized_cc = False
        # Honors the lane lock (a locked cpu-fast box reports/runs "off" even on cross_check=on).
        effective_cc = _manager.effective_cross_check(normalized_cc)

        log["source"] = "url" if url else "file"
        if url:
            log["url_host"] = urlparse(url).hostname or "-"
            # Full URL for debugging. NOTE: presigned URLs put their auth token in
            # the query string, so this lands credentials in the log -- acceptable
            # here (trusted callers, private log); drop the query if that changes.
            log["url"] = url[:512]
        log["seg_model"] = normalized_seg
        log["cross_check"] = "on" if effective_cc else "off"
        log["background_mode"] = effective_bg
        log["screen_color"] = normalized_color if normalized_color else "auto"
        log["auth"] = "ok" if AUTH_TOKEN else "open"

        # Acquire the raw bytes BEFORE the compute gate: a slow fetch must not hold a
        # pool/compute slot. The upload body is already buffered by Starlette (fast);
        # a URL is fetched under its own looser bound (_fetch_sem).
        if url is not None:
            t = time.perf_counter()
            async with _fetch_sem:
                raw = await _fetch_image(url)
            log["t_fetch_ms"] = round((time.perf_counter() - t) * 1000)
        else:
            raw = await file.read()

        # Bound in-flight requests to the pool size, then run on the pool. The GPU
        # sections are serialized inside the pipeline (a shared GPU mutex); the CPU
        # sections of concurrent requests overlap across the pool threads.
        t_wait = time.perf_counter()
        async with _sem:
            log["t_queue_ms"] = round((time.perf_counter() - t_wait) * 1000)
            t_decode = time.perf_counter()
            rgb = _decode_image(raw)
            log["t_decode_ms"] = round((time.perf_counter() - t_decode) * 1000)
            log["input_bytes"] = len(raw)
            log["input_h"], log["input_w"] = int(rgb.shape[0]), int(rgb.shape[1])
            start = time.perf_counter()
            loop = asyncio.get_running_loop()
            try:
                result = await loop.run_in_executor(
                    _executor,
                    lambda: _manager.process(
                        rgb,
                        screen_color=normalized_color,
                        seg_model=normalized_seg,
                        cross_check=normalized_cc,
                        chroma=chroma_enabled,
                    ),
                )
            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")
                log["exc"] = traceback.format_exc()[-2000:]
                raise HTTPException(status_code=500, detail="inference failed")
            except Exception:
                logger.exception("inference failed")
                log["exc"] = traceback.format_exc()[-2000:]
                raise HTTPException(status_code=500, detail="inference failed")
            elapsed_ms = int((time.perf_counter() - start) * 1000)
            log["t_process_ms"] = elapsed_ms

        png = _encode_png(result.rgb, result.alpha)
        log["output_bytes"] = len(png)
        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",
            "X-BGFilter-Background-Mode": effective_bg,
        }
        return Response(content=png, media_type="image/png", headers=headers)
    except HTTPException as exc:
        log["error"] = str(exc.detail)[:200]
        raise