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:
@@ -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)
|
||||
@@ -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()
|
||||
|
||||
@@ -5,10 +5,14 @@ scipy
|
||||
scikit-image
|
||||
pymatting
|
||||
torch
|
||||
torchvision
|
||||
transformers
|
||||
accelerate
|
||||
safetensors
|
||||
huggingface_hub
|
||||
timm
|
||||
einops
|
||||
kornia
|
||||
tqdm
|
||||
typer
|
||||
rich
|
||||
|
||||
Reference in New Issue
Block a user