Spaces:
Running
Running
| """ | |
| bgremove_engine.py — BiRefNet background-removal inference engine. | |
| This module is INDEPENDENT of the LaMa/watermark-removal pipeline and the | |
| Real-ESRGAN upscale pipeline. Do not import or reference any other engine | |
| modules from here. | |
| Architecture | |
| ------------ | |
| BaseBackgroundRemovalEngine(ABC) | |
| └── BiRefNetEngine | |
| ├── __init__: loads BiRefNet via transformers, sets eval mode | |
| └── remove_background: full pipeline | |
| 1. Resize input PIL RGB to 1024×1024 | |
| 2. Normalize with ImageNet mean/std → float32 tensor (1,3,1024,1024) | |
| 3. model(tensor) → preds[-1].sigmoid() → (1,1,1024,1024) mask | |
| 4. Resize mask back to original image dimensions (bilinear, no threshold) | |
| 5. Convert mask to uint8 (0-255) — SOFT mask, NOT thresholded | |
| 6. Composite: original RGB + mask as alpha channel → RGBA PIL image | |
| 7. Return RGBA PIL image | |
| Normalization | |
| ------------- | |
| BiRefNet was trained with standard ImageNet normalization: | |
| mean = [0.485, 0.456, 0.406] | |
| std = [0.229, 0.224, 0.225] | |
| Source confirmed from https://huggingface.co/ZhengPeng7/BiRefNet model card | |
| preprocessing code example and from scratch_test_birefnet_load.py in this project. | |
| Output note | |
| ----------- | |
| The result is RGBA, not RGB. Callers must save as PNG (not JPEG) to preserve | |
| the alpha channel. JPEG does not support transparency. | |
| """ | |
| from abc import ABC, abstractmethod | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import torchvision.transforms as T | |
| from PIL import Image | |
| from transformers import AutoModelForImageSegmentation | |
| from app.config import BIREFNET_MODEL_REPO, BIREFNET_MODEL_CACHE | |
| # Input resolution expected by BiRefNet | |
| _INPUT_SIZE = (1024, 1024) | |
| # ImageNet normalization constants — confirmed from BiRefNet model card | |
| _MEAN = [0.485, 0.456, 0.406] | |
| _STD = [0.229, 0.224, 0.225] | |
| # Alpha cleanup constants — conservative values chosen to reduce fringe/halo | |
| # artifacts at soft edges without sacrificing fine hair/fur detail. | |
| # Erosion of exactly 1 px (3×3 kernel, 1 iteration) removes only the outermost | |
| # contaminated ring; larger values would shave real fur strands. | |
| _ALPHA_ERODE_ITERS = 1 # iterations for cv2.erode (3×3 kernel → ~1 px shrink) | |
| _ALPHA_THRESHOLD = 12 # uint8; soft-suppress alpha values below this toward 0 | |
| _ALPHA_FEATHER_SIGMA = 1.5 # Gaussian sigma (px) applied after erosion to re-soften | |
| # --------------------------------------------------------------------------- | |
| # Alpha mask cleanup helper | |
| # --------------------------------------------------------------------------- | |
| def _cleanup_alpha(alpha: np.ndarray) -> np.ndarray: | |
| """Three-step alpha-channel cleanup to reduce background-colour fringing. | |
| Called AFTER the mask has been resized to original image dimensions, | |
| BEFORE it is composited with the RGB channels. The soft mask is kept | |
| soft throughout — no binary thresholding is applied. | |
| Steps | |
| ----- | |
| 1. Erode (~1 px) | |
| Shrinks the foreground boundary inward by roughly one pixel, removing | |
| the outermost ring of edge pixels. These pixels have the lowest- | |
| confidence (softest) alpha values and are most likely to carry blended | |
| background colour in their RGB channels, producing the visible fringe. | |
| A 3×3 elliptical kernel at 1 iteration gives a conservative 1 px shrink. | |
| 2. Soft threshold (quadratic ramp below _ALPHA_THRESHOLD) | |
| For alpha values below _ALPHA_THRESHOLD (uint8 scale, e.g. 12): | |
| alpha_out = alpha * (alpha / T) | |
| This is a quadratic curve that is 0 at alpha=0 and joins the identity | |
| line continuously at alpha=T (f(T) = T·T/T = T — no cliff). | |
| Effect: a pixel with alpha=6 → 6·(6/12) = 3; alpha=1 → ~0.08 → 0. | |
| Values at or above the threshold are left unchanged. | |
| This pulls faint, nearly-invisible contaminated fringe toward fully | |
| transparent without creating a hard-edged boundary. | |
| 3. Feather (Gaussian blur, sigma=_ALPHA_FEATHER_SIGMA) | |
| Re-softens the boundary that erosion tightened. cv2.GaussianBlur with | |
| (0,0) kernel lets OpenCV derive the kernel size from sigma automatically. | |
| Restores smooth, gradual transitions at fur/hair edges. | |
| Args: | |
| alpha: uint8 numpy array, shape (H, W), values 0–255. | |
| Returns: | |
| uint8 numpy array, same shape, values 0–255. | |
| """ | |
| # 1. Erode — shrink foreground boundary ~1 px | |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) | |
| alpha = cv2.erode(alpha, kernel, iterations=_ALPHA_ERODE_ITERS) | |
| # 2. Soft threshold — quadratic suppression of near-zero fringe alpha values | |
| alpha_f = alpha.astype(np.float32) | |
| low_mask = alpha_f < _ALPHA_THRESHOLD | |
| # f(a) = a * (a / T): 0 at a=0, continuous at a=T, strictly below identity for 0<a<T | |
| alpha_f[low_mask] = alpha_f[low_mask] * (alpha_f[low_mask] / _ALPHA_THRESHOLD) | |
| alpha = np.clip(alpha_f, 0, 255).astype(np.uint8) | |
| # 3. Feather — Gaussian blur to restore soft boundary after erosion | |
| alpha = cv2.GaussianBlur(alpha, (0, 0), sigmaX=_ALPHA_FEATHER_SIGMA) | |
| return alpha | |
| # --------------------------------------------------------------------------- | |
| # Abstract base | |
| # --------------------------------------------------------------------------- | |
| class BaseBackgroundRemovalEngine(ABC): | |
| """Interface for all background removal engines.""" | |
| def remove_background(self, image: Image.Image) -> Image.Image: | |
| """Remove the background from *image*. | |
| Args: | |
| image: Source PIL Image (RGB mode). | |
| Returns: | |
| PIL Image in RGBA mode. The alpha channel carries the soft | |
| foreground mask (255 = fully foreground, 0 = fully background). | |
| The RGB channels are the original pixel colours. | |
| """ | |
| # --------------------------------------------------------------------------- | |
| # BiRefNet implementation | |
| # --------------------------------------------------------------------------- | |
| class BiRefNetEngine(BaseBackgroundRemovalEngine): | |
| """Background remover backed by BiRefNet (ZhengPeng7/BiRefNet). | |
| Loaded once at startup via transformers.AutoModelForImageSegmentation. | |
| Model files are expected to be cached under BIREFNET_MODEL_CACHE | |
| (ml_models/birefnet/) — set at Step 1 of this feature. | |
| """ | |
| def __init__(self) -> None: | |
| """Load BiRefNet from the local cache and set to eval mode. | |
| Raises: | |
| OSError / EnvironmentError: If the model is not found in the cache | |
| and the network is unavailable to download it. | |
| """ | |
| self.model = AutoModelForImageSegmentation.from_pretrained( | |
| BIREFNET_MODEL_REPO, | |
| trust_remote_code=True, | |
| cache_dir=str(BIREFNET_MODEL_CACHE), | |
| ) | |
| self.model = self.model.float() # BiRefNet weights are FP16 on HF Hub; force FP32 for CPU | |
| self.model.eval() | |
| # Preprocessing transform — applied once per call in remove_background | |
| self._transform = T.Compose([ | |
| T.Resize(_INPUT_SIZE), | |
| T.ToTensor(), | |
| T.Normalize(mean=_MEAN, std=_STD), | |
| ]) | |
| # ------------------------------------------------------------------ | |
| # Public API | |
| # ------------------------------------------------------------------ | |
| def remove_background(self, image: Image.Image) -> Image.Image: | |
| """Return an RGBA copy of *image* with the background made transparent. | |
| Pipeline: | |
| 1. Ensure RGB input. | |
| 2. Record original dimensions for mask upscaling. | |
| 3. Preprocess: resize to 1024×1024, ToTensor, ImageNet normalize. | |
| 4. Forward pass through BiRefNet → list of tensors; take last. | |
| 5. Apply sigmoid → float mask in [0, 1], shape (1,1,1024,1024). | |
| 6. Bilinear-upsample mask to original (orig_h, orig_w). | |
| 7. Convert to uint8 [0, 255] — soft mask, no thresholding. | |
| 8. Build RGBA image: original RGB channels + mask as alpha. | |
| 9. Return RGBA PIL Image. | |
| Args: | |
| image: PIL Image (any mode — converted to RGB internally). | |
| Returns: | |
| PIL Image, mode "RGBA", same spatial size as input. | |
| """ | |
| img_rgb = image.convert("RGB") | |
| orig_w, orig_h = img_rgb.size # PIL size is (width, height) | |
| # --- Preprocess: (1, 3, 1024, 1024) float32 tensor ------------------ | |
| input_tensor = self._transform(img_rgb).unsqueeze(0) # (1,3,1024,1024) | |
| # --- Forward pass ----------------------------------------------------- | |
| with torch.no_grad(): | |
| preds = self.model(input_tensor) | |
| # preds is a list/tuple of tensors (one per decoder stage); last is finest | |
| raw_mask = preds[-1].sigmoid() # (1, 1, 1024, 1024), values in [0, 1] | |
| # --- Resize mask back to original image size ------------------------- | |
| # Use bilinear interpolation; align_corners=False is the standard choice | |
| # for image-space operations (avoids corner-pixel alignment artifacts). | |
| mask_resized = F.interpolate( | |
| raw_mask, | |
| size=(orig_h, orig_w), | |
| mode="bilinear", | |
| align_corners=False, | |
| ) # (1, 1, orig_h, orig_w) | |
| # --- Convert to uint8 alpha channel ---------------------------------- | |
| # squeeze: (1,1,H,W) → (H,W) numpy array | |
| mask_np = mask_resized.squeeze().cpu().numpy() # float32, [0, 1] | |
| alpha_u8 = (mask_np * 255.0).clip(0, 255).astype(np.uint8) # uint8, [0, 255] | |
| alpha_u8 = _cleanup_alpha(alpha_u8) # erode + soft-threshold + feather | |
| # --- Composite: RGB + alpha → RGBA ----------------------------------- | |
| rgba_img = img_rgb.copy().convert("RGBA") | |
| alpha_pil = Image.fromarray(alpha_u8, mode="L") | |
| rgba_img.putalpha(alpha_pil) | |
| return rgba_img | |