""" 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 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