lhk229
a6a313d2d0
Adds bgfilter/segmentation.py: BiRefNetSegmenter produces a soft subject mask that defines topology -- it includes colour-contaminated foreground (green-tinted hair) and excludes see-through holes (finger gaps), the distinction chroma/colour cues cannot make. Validated on both samples: green hair -> foreground, finger gap -> background, both correct. Foundation of the semantic-guided redesign; the mask will drive trimap construction next, with ViTMatte still refining soft edges. - SegmentationSettings (ZhengPeng7/BiRefNet, 1024 input, fp32). - Deps: torchvision, timm, einops, kornia. Weights load from HuggingFace via trust_remote_code; behind a firewall set HF_ENDPOINT=https://hf-mirror.com and bypass the proxy. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
74 lines
3.0 KiB
Python
74 lines
3.0 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.no_grad():
|
|
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)
|