lhk229
d24c6ed608
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>
102 lines
3.7 KiB
Python
102 lines
3.7 KiB
Python
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 make_segmenter
|
|
from .settings import PipelineSettings
|
|
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._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)
|
|
|
|
@property
|
|
def default_seg_model(self) -> str:
|
|
return self._base.segmentation.backend
|
|
|
|
@property
|
|
def device(self) -> str:
|
|
return self._base.model.device
|
|
|
|
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 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
|
|
)
|
|
pipeline = MattingPipeline(settings, runner=runner, segmenter=segmenter)
|
|
return pipeline.run_rgb(rgb)
|