Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -12,10 +12,12 @@
 - Fix: Add logging for background preparation issue
 - COMPOSITING FIX: Normalize frame and background scales to prevent dark backgrounds
 - MAJOR FIX: Use enhanced SAM2Handler instead of basic segment_person_hq
+- MAXIMUM QUALITY: Added aggressive mask cleaning for gap elimination
 """
 from __future__ import annotations
 import cv2, numpy as np, os, gc, pickle, logging, tempfile, traceback, threading
 from pathlib import Path
+from typing import Optional, Callable, Dict, Any, Tuple, List
 from .quality_manager import quality_manager # New quality manager import
 
 # Project logger if available
@@ -179,7 +181,7 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
 }
 
 # ---------------------------------------------------------------------------
-# Two-Stage Processor - FIXED VERSION
+# Two-Stage Processor - MAXIMUM QUALITY VERSION
 # ---------------------------------------------------------------------------
 class TwoStageProcessor:
     def __init__(self, sam2_predictor=None, matanyone_model=None):
@@ -213,31 +215,30 @@ def _unwrap_sam2(self, predictor):
         return predictor
 
     def _get_mask(self, frame: np.ndarray) -> np.ndarray:
-        """FIXED: Get segmentation mask using ENHANCED SAM2Handler with full-body detection."""
-        logger.info("=== TwoStageProcessor _get_mask called ===")
+        """MAXIMUM QUALITY mask with enhanced cleaning."""
+        logger.info("=== TwoStageProcessor _get_mask called (MAX QUALITY) ===")
         
         if self.sam2_handler is None:
             logger.warning("No SAM2Handler available - using fallback threshold")
-            # Fallback: simple luminance threshold (kept to avoid breaking callers)
             gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
             _, mask = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
             return mask
         
         try:
-            # CRITICAL FIX: Use the ENHANCED SAM2Handler instead of basic segment_person_hq
             if hasattr(self.sam2_handler, 'create_mask'):
-                logger.info("Using ENHANCED SAM2Handler.create_mask() with full-body detection")
-                # Convert BGR to RGB for SAM2Handler
+                logger.info("Using ENHANCED SAM2Handler.create_mask() with MAXIMUM QUALITY")
                 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 mask = self.sam2_handler.create_mask(frame_rgb)
                 
                 if mask is not None:
-                    logger.info(f"Enhanced SAM2 mask created - shape: {mask.shape}, coverage: {np.mean(mask/255.0):.3f}")
+                    # MAXIMUM QUALITY POST-PROCESSING
+                    mask = self._maximum_quality_mask_cleaning(mask)
+                    logger.info(f"Enhanced SAM2 mask with max quality cleaning - coverage: {np.mean(mask/255.0):.3f}")
                     return mask
                 else:
                     logger.warning("Enhanced SAM2Handler returned None mask")
             else:
-                logger.warning("SAM2Handler doesn't have create_mask method - falling back to basic segmentation")
+                logger.warning("SAM2Handler doesn't have create_mask method")
             
             # Fallback to basic SAM2 if enhanced handler fails
             if self.sam2 is not None:
@@ -245,10 +246,10 @@ def _get_mask(self, frame: np.ndarray) -> np.ndarray:
                 try:
                     from utils.cv_processing import segment_person_hq
                     mask = segment_person_hq(frame, self.sam2)
-                    logger.info(f"Basic SAM2 mask created - coverage: {np.mean(mask/255.0 if mask.max() > 1 else mask):.3f}")
+                    mask = self._maximum_quality_mask_cleaning(mask)
                     return mask
                 except ImportError:
-                    logger.warning("Could not import segment_person_hq - using threshold fallback")
+                    logger.warning("Could not import segment_person_hq")
                 except Exception as e:
                     logger.warning(f"Basic SAM2 segmentation failed: {e}")
             
@@ -261,6 +262,43 @@ def _get_mask(self, frame: np.ndarray) -> np.ndarray:
         _, mask = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
         return mask
 
+    def _maximum_quality_mask_cleaning(self, mask: np.ndarray) -> np.ndarray:
+        """Maximum quality mask cleaning and refinement."""
+        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
+            
+        except Exception as e:
+            logger.warning(f"Maximum quality mask cleaning failed: {e}")
+            return mask
+
     @staticmethod
     def _to_binary_mask(mask: np.ndarray) -> Optional[np.ndarray]:
         """Convert mask to uint8(0..255)."""