Create utils/refinement.py

utils/refinement.py

@@ -0,0 +1,167 @@
+#!/usr/bin/env python3
+"""
+utils.refinement
+─────────────────────────────────────────────────────────────────────────────
+Single-frame mask refinement for BackgroundFX Pro.
+
+Public API
+----------
+refine_mask_hq(image, mask, matanyone_processor, fallback_enabled=True) -> np.ndarray
+"""
+
+from __future__ import annotations
+from typing import Any, Tuple, Optional
+import logging, cv2, torch, numpy as np
+
+log = logging.getLogger(__name__)
+
+# Quality thresholds (same as before)
+MIN_AREA_RATIO = 0.015
+MAX_AREA_RATIO = 0.97
+
+# ────────────────────────────────────────────────────────────────────────────
+# Public
+# ────────────────────────────────────────────────────────────────────────────
+__all__ = ["refine_mask_hq"]
+
+def refine_mask_hq(
+    image: np.ndarray,
+    mask:  np.ndarray,
+    matanyone_processor: Any,
+    fallback_enabled: bool = True,
+) -> np.ndarray:
+    """
+    1) Try MatAnyOne high-quality refinement.
+    2) Otherwise OpenCV “enhanced” filter.
+    3) GrabCut and saliency fallbacks.
+    Always returns uint8 mask (0/255).
+    """
+    mask = _process_mask(mask)
+
+    # 1 — MatAnyOne
+    if matanyone_processor is not None:
+        try:
+            refined = _matanyone_refine(image, mask, matanyone_processor)
+            if refined is not None and _validate_mask_quality(refined, image.shape[:2]):
+                return refined
+            log.warning("MatAnyOne produced poor mask; fallback")
+        except Exception as e:
+            log.warning(f"MatAnyOne error: {e}")
+
+    # 2 — OpenCV “enhanced” bilateral+guided+MORPH
+    try:
+        refined = _opencv_enhance(image, mask)
+        if _validate_mask_quality(refined, image.shape[:2]):
+            return refined
+    except Exception as e:
+        log.debug(f"OpenCV enhance error: {e}")
+
+    # 3 — GrabCut + saliency double-fallback
+    try:
+        gc = _refine_with_grabcut(image, mask)
+        if _validate_mask_quality(gc, image.shape[:2]):
+            return gc
+        sal = _refine_with_saliency(image, mask)
+        if _validate_mask_quality(sal, image.shape[:2]):
+            return sal
+    except Exception as e:
+        log.debug(f"GrabCut/saliency fallback error: {e}")
+
+    # last resort
+    return mask if fallback_enabled else _opencv_enhance(image, mask)
+
+# ────────────────────────────────────────────────────────────────────────────
+# MatAnyOne wrapper (safe)
+# ────────────────────────────────────────────────────────────────────────────
+def _matanyone_refine(img, mask, proc) -> Optional[np.ndarray]:
+    if not (hasattr(proc, "step") and hasattr(proc, "output_prob_to_mask")):
+        return None
+    # image tensor (C,H,W) float32 0-1
+    anp = img.astype(np.float32)
+    if anp.max() > 1: anp /= 255.0
+    anp = np.transpose(anp, (2,0,1))
+    img_t  = torch.from_numpy(anp).unsqueeze(0).to(proc.device if hasattr(proc,"device") else "cpu")
+    mask_f = mask.astype(np.float32)/255.0
+    mask_t = torch.from_numpy(mask_f).unsqueeze(0).to(img_t.device)
+
+    with torch.no_grad():
+        prob = proc.step(img_t, mask_t, objects=[1])
+        m = proc.output_prob_to_mask(prob).squeeze().cpu().numpy()
+    if m.max() <= 1: m *= 255
+    return m.astype(np.uint8)
+
+# ────────────────────────────────────────────────────────────────────────────
+# OpenCV enhanced filter chain
+# ────────────────────────────────────────────────────────────────────────────
+def _opencv_enhance(img, mask):
+    if mask.ndim == 3: mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+    if mask.max()<=1:  mask = (mask*255).astype(np.uint8)
+    m = cv2.bilateralFilter(mask, 9, 75, 75)
+    m = _guided_filter(img, m, r=8, eps=0.2)
+    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5)))
+    m = cv2.morphologyEx(m, cv2.MORPH_OPEN,  cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)))
+    m = cv2.GaussianBlur(m,(3,3),0.8)
+    _,m = cv2.threshold(m,127,255,cv2.THRESH_BINARY)
+    return m
+
+def _guided_filter(guide, mask, r=8, eps=0.2):
+    g = cv2.cvtColor(guide, cv2.COLOR_BGR2GRAY).astype(np.float32)/255.0
+    m = mask.astype(np.float32)/255.0
+    k = 2*r+1
+    mean_g  = cv2.boxFilter(g, -1, (k,k))
+    mean_m  = cv2.boxFilter(m, -1, (k,k))
+    corr_gm = cv2.boxFilter(g*m, -1, (k,k))
+    cov     = corr_gm - mean_g*mean_m
+    var_g   = cv2.boxFilter(g*g, -1, (k,k)) - mean_g*mean_g
+    a = cov/(var_g+eps)
+    b = mean_m - a*mean_g
+    mean_a = cv2.boxFilter(a, -1, (k,k))
+    mean_b = cv2.boxFilter(b, -1, (k,k))
+    out = (mean_a*g+mean_b)*255
+    return out.astype(np.uint8)
+
+# ────────────────────────────────────────────────────────────────────────────
+# GrabCut & saliency fallbacks
+# ────────────────────────────────────────────────────────────────────────────
+def _refine_with_grabcut(img, seed):
+    h,w = img.shape[:2]
+    gc = np.full((h,w), cv2.GC_PR_BGD, np.uint8)
+    gc[seed>200] = cv2.GC_FGD
+    rect = (w//4, h//6, w//2, int(h*0.7))
+    bgd,fgd = np.zeros((1,65),np.float64), np.zeros((1,65),np.float64)
+    cv2.grabCut(img, gc, rect, bgd, fgd, 3, cv2.GC_INIT_WITH_MASK)
+    return np.where((gc==cv2.GC_FGD)|(gc==cv2.GC_PR_FGD),255,0).astype(np.uint8)
+
+def _refine_with_saliency(img, seed):
+    sal = _compute_saliency(img)
+    if sal is None: return seed
+    high = (sal>0.6).astype(np.uint8)*255
+    cy,cx = img.shape[0]//2, img.shape[1]//2
+    if np.any(seed>127):
+        ys,xs = np.where(seed>127); cy,cx=int(np.mean(ys)),int(np.mean(xs))
+    ff = high.copy(); cv2.floodFill(ff,None,(cx,cy),255,loDiff=5,upDiff=5)
+    return ff
+
+def _compute_saliency(img):
+    try:
+        if hasattr(cv2,"saliency"):
+            s=cv2.saliency.StaticSaliencySpectralResidual_create()
+            ok,sm=s.computeSaliency(img)
+            if ok: return (sm-sm.min())/max(1e-6,sm.max()-sm.min())
+    except Exception: pass
+    return None
+
+# ────────────────────────────────────────────────────────────────────────────
+# Helpers
+# ────────────────────────────────────────────────────────────────────────────
+def _process_mask(mask):
+    if mask.ndim==3: mask=cv2.cvtColor(mask,cv2.COLOR_BGR2GRAY)
+    if mask.dtype!=np.uint8:
+        mask = (mask*255).astype(np.uint8) if mask.max()<=1 else mask.astype(np.uint8)
+    _,mask=cv2.threshold(mask,127,255,cv2.THRESH_BINARY)
+    return mask
+
+def _validate_mask_quality(mask, shape: Tuple[int,int]) -> bool:
+    h,w = shape
+    ratio = np.sum(mask>127)/(h*w)
+    return MIN_AREA_RATIO <= ratio <= MAX_AREA_RATIO