From a6a313d2d0ce0832dc0ffab90858a7c3764060cb Mon Sep 17 00:00:00 2001
From: lingh <ink29535@proton.me>
Date: Tue, 30 Jun 2026 21:07:24 +0800
Subject: [PATCH] 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>
---
 bgfilter/segmentation.py | 73 ++++++++++++++++++++++++++++++++++++++++
 bgfilter/settings.py     |  8 +++++
 requirements.txt         |  4 +++
 3 files changed, 85 insertions(+)
 create mode 100644 bgfilter/segmentation.py

diff --git a/bgfilter/segmentation.py b/bgfilter/segmentation.py
new file mode 100644
index 0000000..4328d11
--- /dev/null
+++ b/bgfilter/segmentation.py
@@ -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)
diff --git a/bgfilter/settings.py b/bgfilter/settings.py
index 9fb2c30..81c9945 100644
--- a/bgfilter/settings.py
+++ b/bgfilter/settings.py
@@ -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()
diff --git a/requirements.txt b/requirements.txt
index c48f6f6..6fddb5d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,10 +5,14 @@ scipy
 scikit-image
 pymatting
 torch
+torchvision
 transformers
 accelerate
 safetensors
 huggingface_hub
+timm
+einops
+kornia
 tqdm
 typer
 rich