Add BiRefNet semantic segmentation module

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>
This commit is contained in:
2026-06-30 21:07:24 +08:00
parent a9886f1479
commit a6a313d2d0
3 changed files with 85 additions and 0 deletions
+73
View File
@@ -0,0 +1,73 @@
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)
+8
View File
@@ -78,6 +78,14 @@ class ModelSettings:
fallback_to_chroma_alpha: bool = False
@dataclass(frozen=True)
class SegmentationSettings:
enabled: bool = True
model_name: str = "ZhengPeng7/BiRefNet"
device: str = "cuda"
input_size: int = 1024
@dataclass(frozen=True)
class PipelineSettings:
chroma: ChromaSettings = ChromaSettings()
+4
View File
@@ -5,10 +5,14 @@ scipy
scikit-image
pymatting
torch
torchvision
transformers
accelerate
safetensors
huggingface_hub
timm
einops
kornia
tqdm
typer
rich