BGfilter/bgfilter/service.py

from __future__ import annotations

import logging
import threading
from dataclasses import replace

import numpy as np

from .config import _SEG_MODELS, override_settings
from .pipeline import MattingPipeline, MattingResult
from .segmentation import BiRefNetSegmenter, make_segmenter
from .settings import PipelineSettings, SegmentationSettings
from .vitmatte_infer import ViTMatteRunner

logger = logging.getLogger("bgfilter.service")

# HTTP-facing seg_model values, in preference order (first = default).
SEG_MODELS: tuple[str, ...] = tuple(_SEG_MODELS)


class PipelineManager:
    """Long-lived holder of the matting models.

    ViTMatte is the matting model shared by both segmentation backends, so it is
    loaded once and reused. Each segmentation backend is lazily loaded once on
    first use and cached by name. Per-request parameters (``screen_color``,
    ``seg_model``) never trigger a model reload -- ``screen_color`` only affects
    this request's chroma/despill, and ``seg_model`` just selects an already-cached
    (or first-time-loaded) segmenter.
    """

    def __init__(self, base_settings: PipelineSettings, preload: bool = True):
        self._base = base_settings
        self._runner: ViTMatteRunner | None = None
        self._segmenters: dict[str, object] = {}
        self._cross_checker: BiRefNetSegmenter | None = None
        self._load_lock = threading.Lock()
        if preload:
            if base_settings.model.matting_method == "vitmatte":
                self._get_runner()
            if base_settings.segmentation.enabled:
                self._get_segmenter(base_settings.segmentation.backend)
            if base_settings.cross_check.enabled:
                self._get_cross_checker()

    @property
    def default_seg_model(self) -> str:
        return self._base.segmentation.backend

    @property
    def device(self) -> str:
        return self._base.model.device

    @property
    def cross_check_enabled(self) -> bool:
        return self._base.cross_check.enabled

    def _get_runner(self) -> ViTMatteRunner:
        if self._runner is None:
            with self._load_lock:
                if self._runner is None:
                    logger.info(
                        "loading ViTMatte: model=%s device=%s",
                        self._base.model.model_name,
                        self._base.model.device,
                    )
                    self._runner = ViTMatteRunner(self._base.model)
        return self._runner

    def _get_segmenter(self, backend: str):
        seg = self._segmenters.get(backend)
        if seg is None:
            with self._load_lock:
                seg = self._segmenters.get(backend)
                if seg is None:
                    seg_settings = replace(
                        self._base.segmentation,
                        backend=backend,
                        model_name=_SEG_MODELS[backend],
                    )
                    logger.info(
                        "loading segmenter: backend=%s model=%s device=%s",
                        backend,
                        seg_settings.model_name,
                        seg_settings.device,
                    )
                    seg = make_segmenter(seg_settings)
                    self._segmenters[backend] = seg
        return seg

    def _get_cross_checker(self) -> BiRefNetSegmenter:
        # The cross-check model depends only on base settings (never on per-request
        # params), so one shared instance serves all requests.
        if self._cross_checker is None:
            with self._load_lock:
                if self._cross_checker is None:
                    check = self._base.cross_check
                    logger.info(
                        "loading cross-check model: model=%s device=%s",
                        check.model_name,
                        self._base.segmentation.device,
                    )
                    self._cross_checker = BiRefNetSegmenter(
                        SegmentationSettings(
                            model_name=check.model_name,
                            device=self._base.segmentation.device,
                            input_size=check.input_size,
                        )
                    )
        return self._cross_checker

    def process(
        self, rgb: np.ndarray, *, screen_color: str | None, seg_model: str
    ) -> MattingResult:
        if seg_model not in _SEG_MODELS:
            raise ValueError(f"invalid seg_model '{seg_model}'")
        settings = override_settings(
            self._base, screen_color=screen_color, seg_backend=seg_model
        )
        runner = (
            self._get_runner()
            if settings.model.matting_method == "vitmatte"
            else None
        )
        segmenter = (
            self._get_segmenter(seg_model) if settings.segmentation.enabled else None
        )
        cross_checker = (
            self._get_cross_checker() if settings.cross_check.enabled else None
        )
        pipeline = MattingPipeline(
            settings, runner=runner, segmenter=segmenter, cross_checker=cross_checker
        )
        return pipeline.run_rgb(rgb)