BGfilter/bgfilter/segmentation.py

from __future__ import annotations

import os

import numpy as np
from PIL import Image

from .precision import ensure_deform_conv2d_autocast_cpu, resolve_dtype
from .settings import SegmentationSettings
from .weights import resolve_model_source


class BiRefNetSegmenter:
    """Semantic subject segmentation via BiRefNet.

    Produces a soft foreground mask that defines subject *topology*: it includes
    colour-contaminated regions (e.g. green-tinted hair) as foreground and
    excludes see-through holes (e.g. gaps between fingers) as background -- the
    distinction that chroma/colour cues cannot make. The mask drives trimap
    construction; ViTMatte still refines the soft edges.

    Environment notes:
    - Weights load from HuggingFace. Behind a firewall set ``HF_ENDPOINT``
      (e.g. ``https://hf-mirror.com``); a flaky proxy may need to be bypassed.
    - The repo ships custom modeling code, so ``trust_remote_code=True`` is
      required (it executes that code -- only use a model source you trust).
    """

    def __init__(self, settings: SegmentationSettings):
        try:
            import torch
            from torchvision import transforms
            from transformers import AutoModelForImageSegmentation
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "Missing segmentation dependencies. Install them with: "
                "pip install torch torchvision transformers timm einops kornia"
            ) from exc

        self.torch = torch
        self.model = AutoModelForImageSegmentation.from_pretrained(
            resolve_model_source(settings.model_name), trust_remote_code=True
        )
        self.model.eval()
        self.model.float()  # checkpoint ships as fp16; force fp32 to match inputs
        self.device = self._resolve_device(settings.device)
        self.model.to(self.device)
        # bf16 runs via autocast (weights stay fp32): conv/matmul take the fast
        # bf16 kernels while ops without them keep fp32. deform_conv2d has no
        # bf16 CPU kernel and no AutocastCPU wrapper in torchvision, so it needs
        # a dispatcher-level fp32 shim or the forward crashes.
        self.dtype = resolve_dtype(torch, self.device, settings.precision)
        if self.dtype == torch.bfloat16 and self.device.type == "cpu":
            ensure_deform_conv2d_autocast_cpu(torch)

        size = settings.input_size
        self.transform = transforms.Compose(
            [
                transforms.Resize((size, size)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ]
        )

    def _resolve_device(self, requested: str):
        if requested == "cuda" and not self.torch.cuda.is_available():
            raise RuntimeError(
                "CUDA was requested but is not available. Use --device cpu or install a CUDA-enabled PyTorch build."
            )
        return self.torch.device(requested)

    def mask(self, rgb: np.ndarray) -> np.ndarray:
        """Return a soft foreground mask, float32 0..1, at the input resolution."""
        image = Image.fromarray(rgb.astype(np.uint8), mode="RGB")
        tensor = self.transform(image).unsqueeze(0).to(self.device)
        autocast = self.torch.autocast(
            self.device.type,
            dtype=self.torch.bfloat16,
            enabled=self.dtype == self.torch.bfloat16,
        )
        with self.torch.inference_mode(), autocast:
            pred = self.model(tensor)[-1].sigmoid()
        pred = pred[0, 0].detach().float().cpu().numpy()
        if pred.shape != rgb.shape[:2]:
            pred_img = Image.fromarray(
                np.clip(pred * 255.0, 0, 255).astype(np.uint8), mode="L"
            ).resize((rgb.shape[1], rgb.shape[0]), Image.Resampling.BILINEAR)
            pred = np.asarray(pred_img, dtype=np.float32) / 255.0
        return np.clip(pred, 0.0, 1.0).astype(np.float32)


class AnimeSegSegmenter:
    """Anime character segmentation via SkyTNT anime-seg (ISNet, ONNX).

    Anime-trained, run through onnxruntime (no remote code). Same soft-mask
    interface as :class:`BiRefNetSegmenter`. The ONNX output is already 0..1.
    Weights load from HuggingFace; behind a firewall set ``HF_ENDPOINT``.
    """

    def __init__(self, settings: SegmentationSettings):
        try:
            import onnxruntime as ort
            from huggingface_hub import hf_hub_download
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "Missing anime-seg dependencies. Install them with: "
                "pip install onnxruntime huggingface_hub"
            ) from exc

        source = resolve_model_source(settings.model_name)
        if os.path.isdir(source):
            model_file = os.path.join(source, "isnetis.onnx")
        else:
            model_file = hf_hub_download(source, "isnetis.onnx")
        providers = ["CPUExecutionProvider"]
        if settings.device == "cuda" and "CUDAExecutionProvider" in ort.get_available_providers():
            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self.session = ort.InferenceSession(model_file, providers=providers)
        self.input_name = self.session.get_inputs()[0].name
        self.size = settings.input_size

    def mask(self, rgb: np.ndarray) -> np.ndarray:
        image = Image.fromarray(rgb.astype(np.uint8), mode="RGB").resize(
            (self.size, self.size), Image.Resampling.BILINEAR
        )
        x = (np.asarray(image, dtype=np.float32) / 255.0).transpose(2, 0, 1)[None]
        pred = self.session.run(None, {self.input_name: x})[0][0, 0]
        mask_img = Image.fromarray(
            np.clip(pred * 255.0, 0, 255).astype(np.uint8), mode="L"
        ).resize((rgb.shape[1], rgb.shape[0]), Image.Resampling.BILINEAR)
        return np.asarray(mask_img, dtype=np.float32) / 255.0


def make_segmenter(settings: SegmentationSettings):
    if settings.backend == "birefnet":
        return BiRefNetSegmenter(settings)
    if settings.backend == "anime-seg":
        return AnimeSegSegmenter(settings)
    raise RuntimeError(
        f"Unknown segmentation backend '{settings.backend}'. Use 'birefnet' or 'anime-seg'."
    )