Update utils/cv_processing.py

utils/cv_processing.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 """
-cv_processing.py · FIXED VERSION with proper SAM2 handling + MatAnyone stateful integration
-Now with environment variable support for USE_SAM2 and USE_MATANYONE
+cv_processing.py · MAXIMUM QUALITY VERSION with enhanced SAM2Handler integration
+Updated to work with enhanced SAM2Handler that has full-body detection strategies
+Now includes maximum quality mask cleaning and aggressive post-processing
 
 All public functions in this module expect RGB images (H,W,3) unless stated otherwise.
 CoreVideoProcessor already converts BGR→RGB before calling into this module.
@@ -32,6 +33,11 @@ def _use_matanyone_enabled() -> bool:
     val = os.getenv("USE_MATANYONE", "1")
     return val.lower() in ("1", "true", "yes", "on")
 
+def _use_max_quality_enabled() -> bool:
+    """Check if maximum quality processing should be used"""
+    val = os.getenv("BFX_QUALITY", "max")
+    return val.lower() == "max"
+
 # ----------------------------------------------------------------------------
 # Background presets
 # ----------------------------------------------------------------------------
@@ -139,6 +145,46 @@ def _vertical_gradient(top: Tuple[int,int,int], bottom: Tuple[int,int,int], widt
         bg[y, :] = (r, g, b)
     return bg
 
+# ----------------------------------------------------------------------------
+# Maximum Quality Mask Cleaning (integrated from TwoStageProcessor)
+# ----------------------------------------------------------------------------
+def _maximum_quality_mask_cleaning(mask: np.ndarray) -> np.ndarray:
+    """Maximum quality mask cleaning and refinement - same as TwoStageProcessor."""
+    try:
+        # Ensure uint8 format
+        if mask.max() <= 1.0:
+            mask_uint8 = (mask * 255).astype(np.uint8)
+        else:
+            mask_uint8 = mask.astype(np.uint8)
+        
+        # Step 1: Fill small holes aggressively
+        kernel_fill = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
+        mask_filled = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel_fill)
+        
+        # Step 2: Connect nearby regions
+        kernel_connect = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
+        mask_connected = cv2.morphologyEx(mask_filled, cv2.MORPH_CLOSE, kernel_connect)
+        
+        # Step 3: Smooth boundaries heavily
+        mask_smooth1 = cv2.GaussianBlur(mask_connected, (7, 7), 2.0)
+        
+        # Step 4: Re-threshold to crisp edges
+        _, mask_thresh = cv2.threshold(mask_smooth1, 127, 255, cv2.THRESH_BINARY)
+        
+        # Step 5: Final edge smoothing
+        mask_final = cv2.GaussianBlur(mask_thresh, (5, 5), 1.0)
+        
+        # Step 6: Dilate slightly to ensure full coverage
+        kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        mask_dilated = cv2.dilate(mask_final, kernel_dilate, iterations=1)
+        
+        logger.info("Maximum quality mask cleaning applied successfully")
+        return (mask_dilated.astype(np.float32) / 255.0)
+        
+    except Exception as e:
+        logger.warning(f"Maximum quality mask cleaning failed: {e}")
+        return mask
+
 # ----------------------------------------------------------------------------
 # Background creation
 # ----------------------------------------------------------------------------
@@ -160,7 +206,7 @@ def create_professional_background(key_or_cfg: Any, width: int, height: int) ->
     return _vertical_gradient(dark, color, width, height)
 
 # ----------------------------------------------------------------------------
-# Improved Segmentation (expects RGB input)
+# Improved Segmentation (expects RGB input) - ENHANCED FOR SAM2Handler
 # ----------------------------------------------------------------------------
 def _simple_person_segmentation(frame_rgb: np.ndarray) -> np.ndarray:
     """Basic fallback segmentation using color detection on RGB frames."""
@@ -189,7 +235,14 @@ def _simple_person_segmentation(frame_rgb: np.ndarray) -> np.ndarray:
         person_mask = np.zeros_like(person_mask)
         cv2.drawContours(person_mask, [largest_contour], -1, 255, -1)
 
-    return (person_mask.astype(np.float32) / 255.0)
+    mask_result = (person_mask.astype(np.float32) / 255.0)
+    
+    # Apply maximum quality cleaning if enabled
+    if _use_max_quality_enabled():
+        mask_result = _maximum_quality_mask_cleaning(mask_result)
+        logger.info("Applied maximum quality cleaning to fallback segmentation")
+    
+    return mask_result
 
 def segment_person_hq(
     frame: np.ndarray,
@@ -199,7 +252,8 @@ def segment_person_hq(
     **_compat_kwargs,
 ) -> np.ndarray:
     """
-    High-quality person segmentation with proper SAM2 handling.
+    High-quality person segmentation with ENHANCED SAM2Handler integration.
+    Now uses enhanced SAM2Handler.create_mask() for full-body detection.
     Expects RGB frame (H,W,3), uint8 or float in [0,1].
     """
     # Override with environment variable if not explicitly set
@@ -215,7 +269,38 @@ def segment_person_hq(
 
     if predictor is not None:
         try:
-            if hasattr(predictor, "set_image") and hasattr(predictor, "predict"):
+            # ENHANCED: Check if this is the new SAM2Handler with create_mask method
+            if hasattr(predictor, 'create_mask'):
+                logger.info("Using ENHANCED SAM2Handler.create_mask() with full-body detection")
+                # SAM2Handler expects RGB uint8
+                if frame_rgb.dtype != np.uint8:
+                    rgb_u8 = np.clip(frame_rgb * (255.0 if frame_rgb.dtype != np.uint8 else 1.0), 0, 255).astype(np.uint8) \
+                             if np.issubdtype(frame_rgb.dtype, np.floating) else frame_rgb.astype(np.uint8)
+                else:
+                    rgb_u8 = frame_rgb
+                
+                # Use enhanced SAM2Handler with full-body detection strategies
+                mask = predictor.create_mask(rgb_u8)
+                
+                if mask is not None:
+                    # Convert to float format
+                    mask_float = _to_mask01(mask)
+                    logger.info(f"Enhanced SAM2Handler mask stats: shape={mask_float.shape}, min={mask_float.min():.3f}, max={mask_float.max():.3f}, mean={mask_float.mean():.3f}")
+                    
+                    if float(mask_float.max()) > 0.1:
+                        # Apply additional maximum quality cleaning if enabled
+                        if _use_max_quality_enabled():
+                            mask_float = _maximum_quality_mask_cleaning(mask_float)
+                            logger.info("Applied additional maximum quality cleaning to enhanced SAM2 result")
+                        return np.ascontiguousarray(mask_float)
+                    else:
+                        logger.warning("Enhanced SAM2Handler mask too weak, using fallback")
+                else:
+                    logger.warning("Enhanced SAM2Handler returned None mask")
+            
+            # FALLBACK: Basic SAM2 predictor handling (legacy compatibility)
+            elif hasattr(predictor, "set_image") and hasattr(predictor, "predict"):
+                logger.info("Using legacy SAM2 predictor interface")
                 # Predictor adapter expects RGB uint8; convert if needed
                 if frame_rgb.dtype != np.uint8:
                     rgb_u8 = np.clip(frame_rgb * (255.0 if frame_rgb.dtype != np.uint8 else 1.0), 0, 255).astype(np.uint8) \
@@ -272,13 +357,20 @@ def segment_person_hq(
                     if mask is not None:
                         mask = _to_mask01(mask)
                         # Add debug logging
-                        logger.info(f"SAM2 mask stats: shape={mask.shape}, min={mask.min():.3f}, max={mask.max():.3f}, mean={mask.mean():.3f}")
+                        logger.info(f"Legacy SAM2 mask stats: shape={mask.shape}, min={mask.min():.3f}, max={mask.max():.3f}, mean={mask.mean():.3f}")
+                        
                         if float(mask.max()) > 0.1:
+                            # Apply maximum quality cleaning if enabled
+                            if _use_max_quality_enabled():
+                                mask = _maximum_quality_mask_cleaning(mask)
+                                logger.info("Applied maximum quality cleaning to legacy SAM2 result")
                             return np.ascontiguousarray(mask)
                         else:
-                            logger.warning("SAM2 mask too weak, using fallback")
-                else:
-                    logger.warning("SAM2 returned no masks")
+                            logger.warning("Legacy SAM2 mask too weak, using fallback")
+                    else:
+                        logger.warning("Legacy SAM2 returned no masks")
+            else:
+                logger.warning("Predictor doesn't have expected SAM2 interface")
 
         except Exception as e:
             logger.warning(f"SAM2 segmentation error: {e}")
@@ -292,7 +384,7 @@ def segment_person_hq(
 segment_person_hq_original = segment_person_hq
 
 # ----------------------------------------------------------------------------
-# MatAnyone Refinement (Stateful-capable)
+# MatAnyone Refinement (Stateful-capable) - ENHANCED WITH MAX QUALITY
 # ----------------------------------------------------------------------------
 def refine_mask_hq(
     frame: np.ndarray,
@@ -305,7 +397,7 @@ def refine_mask_hq(
     **_compat_kwargs,
 ) -> np.ndarray:
     """
-    Refine mask with MatAnyone.
+    Refine mask with MatAnyone + maximum quality post-processing.
 
     Modes:
       • Stateful (preferred): provide `frame_idx`. On frame_idx==0, the session encodes with the mask.
@@ -324,6 +416,10 @@ def refine_mask_hq(
 
     if use_matanyone is False:
         logger.info("MatAnyone disabled by environment variable, returning unrefined mask")
+        # Still apply maximum quality cleaning if enabled
+        if _use_max_quality_enabled():
+            mask01 = _maximum_quality_mask_cleaning(mask01)
+            logger.info("Applied maximum quality cleaning to unrefined mask")
         return mask01
 
     if matanyone is not None and callable(matanyone):
@@ -338,7 +434,8 @@ def refine_mask_hq(
                     refined = matanyone(rgb01)                 # propagate without mask
                 refined = _mask_to_2d(refined)
                 if float(refined.max()) > 0.1:
-                    return _postprocess_mask(refined)
+                    result = _postprocess_mask_max_quality(refined)
+                    return result
                 logger.warning("MatAnyone stateful refinement produced empty/weak mask; falling back.")
 
             # Backward-compat (stateless) path
@@ -368,7 +465,8 @@ def refine_mask_hq(
                     logger.debug(f"MatAnyone process failed: {e}")
 
             if refined is not None and float(refined.max()) > 0.1:
-                return _postprocess_mask(refined)
+                result = _postprocess_mask_max_quality(refined)
+                return result
             else:
                 logger.warning("MatAnyone refinement failed or produced empty mask")
 
@@ -377,12 +475,27 @@ def refine_mask_hq(
 
     # Fallback refinement
     if fallback_enabled:
-        return _fallback_refine(mask01)
+        return _fallback_refine_max_quality(mask01)
     else:
+        # Still apply maximum quality cleaning if enabled
+        if _use_max_quality_enabled():
+            mask01 = _maximum_quality_mask_cleaning(mask01)
+            logger.info("Applied maximum quality cleaning to fallback mask")
         return mask01
 
+def _postprocess_mask_max_quality(mask01: np.ndarray) -> np.ndarray:
+    """Post-process mask with maximum quality cleaning"""
+    if _use_max_quality_enabled():
+        # Use the aggressive maximum quality cleaning
+        result = _maximum_quality_mask_cleaning(mask01)
+        logger.info("Applied maximum quality post-processing to MatAnyone result")
+        return result
+    else:
+        # Use standard post-processing
+        return _postprocess_mask(mask01)
+
 def _postprocess_mask(mask01: np.ndarray) -> np.ndarray:
-    """Post-process mask to clean edges and remove artifacts"""
+    """Standard post-process mask to clean edges and remove artifacts"""
     mask_uint8 = (np.clip(mask01, 0, 1) * 255).astype(np.uint8)
 
     kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
@@ -397,6 +510,17 @@ def _postprocess_mask(mask01: np.ndarray) -> np.ndarray:
     out = mask_uint8.astype(np.float32) / 255.0
     return np.ascontiguousarray(out)
 
+def _fallback_refine_max_quality(mask01: np.ndarray) -> np.ndarray:
+    """Fallback refinement with maximum quality option"""
+    if _use_max_quality_enabled():
+        # Use aggressive maximum quality cleaning
+        result = _maximum_quality_mask_cleaning(mask01)
+        logger.info("Applied maximum quality cleaning to fallback refinement")
+        return result
+    else:
+        # Use standard fallback refinement
+        return _fallback_refine(mask01)
+
 def _fallback_refine(mask01: np.ndarray) -> np.ndarray:
     """Simple fallback refinement"""
     mask_uint8 = (np.clip(mask01, 0, 1) * 255).astype(np.uint8)
@@ -413,7 +537,7 @@ def _fallback_refine(mask01: np.ndarray) -> np.ndarray:
     return np.ascontiguousarray(out)
 
 # ----------------------------------------------------------------------------
-# Compositing (expects RGB inputs)
+# Compositing (expects RGB inputs) - ENHANCED WITH MAX QUALITY
 # ----------------------------------------------------------------------------
 def replace_background_hq(
     frame: np.ndarray,
@@ -422,7 +546,7 @@ def replace_background_hq(
     fallback_enabled: bool = True,
     **_compat,
 ) -> np.ndarray:
-    """High-quality background replacement with alpha blending (RGB in/out)."""
+    """High-quality background replacement with alpha blending (RGB in/out) - enhanced with max quality."""
     try:
         H, W = frame.shape[:2]
 
@@ -431,7 +555,14 @@ def replace_background_hq(
 
         m = _mask_to_2d(_to_mask01(mask01))
 
-        m = _feather(m, k=1)
+        # Apply maximum quality cleaning to mask before compositing
+        if _use_max_quality_enabled():
+            m = _maximum_quality_mask_cleaning(m)
+            logger.debug("Applied maximum quality cleaning to compositing mask")
+
+        # Enhanced feathering for maximum quality
+        feather_strength = 3 if _use_max_quality_enabled() else 1
+        m = _feather(m, k=feather_strength)
 
         m3 = np.repeat(m[:, :, None], 3, axis=2)