lhk229
8b574bc551
Both call sites convert outputs to numpy immediately, so the stricter inference-mode tensors are safe; saves autograd view/version tracking overhead. Verified end-to-end on CUDA (alpha finite, no NaN). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
122 lines
5.1 KiB
Python
122 lines
5.1 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from .settings import SegmentationSettings
|
|
|
|
|
|
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(
|
|
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)
|
|
|
|
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)
|
|
with self.torch.inference_mode():
|
|
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
|
|
|
|
model_file = hf_hub_download(settings.model_name, "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'."
|
|
)
|