Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -12,7 +12,7 @@
 - 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
+- TESTING FIX: Disabled double mask cleaning to eliminate artifacts
 """
 from __future__ import annotations
 import cv2, numpy as np, os, gc, pickle, logging, tempfile, traceback, threading
@@ -181,7 +181,7 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
 }
 
 # ---------------------------------------------------------------------------
-# Two-Stage Processor - MAXIMUM QUALITY VERSION
+# Two-Stage Processor - TESTING: NO DOUBLE CLEANING
 # ---------------------------------------------------------------------------
 class TwoStageProcessor:
     def __init__(self, sam2_predictor=None, matanyone_model=None):
@@ -215,8 +215,8 @@ def _unwrap_sam2(self, predictor):
         return predictor
 
     def _get_mask(self, frame: np.ndarray) -> np.ndarray:
-        """MAXIMUM QUALITY mask with enhanced cleaning."""
-        logger.info("=== TwoStageProcessor _get_mask called (MAX QUALITY) ===")
+        """TESTING: Disable double cleaning to isolate mask quality issues."""
+        logger.info("=== TwoStageProcessor _get_mask called (TESTING - NO DOUBLE CLEANING) ===")
         
         if self.sam2_handler is None:
             logger.warning("No SAM2Handler available - using fallback threshold")
@@ -226,17 +226,17 @@ def _get_mask(self, frame: np.ndarray) -> np.ndarray:
         
         try:
             if hasattr(self.sam2_handler, 'create_mask'):
-                logger.info("Using ENHANCED SAM2Handler.create_mask() with MAXIMUM QUALITY")
+                logger.info("Using SAM2Handler.create_mask() WITHOUT additional cleaning")
                 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 mask = self.sam2_handler.create_mask(frame_rgb)
                 
                 if mask is not None:
-                    # 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}")
+                    # DISABLED: No additional cleaning - let SAM2Handler's cleaning stand
+                    # mask = self._maximum_quality_mask_cleaning(mask)  # DISABLED FOR TESTING
+                    logger.info(f"SAM2 mask WITHOUT additional cleaning - coverage: {np.mean(mask/255.0):.3f}")
                     return mask
                 else:
-                    logger.warning("Enhanced SAM2Handler returned None mask")
+                    logger.warning("SAM2Handler returned None mask")
             else:
                 logger.warning("SAM2Handler doesn't have create_mask method")
             
@@ -246,7 +246,8 @@ 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)
-                    mask = self._maximum_quality_mask_cleaning(mask)
+                    # DISABLED: No additional cleaning in fallback either
+                    # mask = self._maximum_quality_mask_cleaning(mask)  # DISABLED FOR TESTING
                     return mask
                 except ImportError:
                     logger.warning("Could not import segment_person_hq")
@@ -254,7 +255,7 @@ def _get_mask(self, frame: np.ndarray) -> np.ndarray:
                     logger.warning(f"Basic SAM2 segmentation failed: {e}")
             
         except Exception as e:
-            logger.error(f"Enhanced SAM2 segmentation failed: {e}")
+            logger.error(f"SAM2 segmentation failed: {e}")
         
         # Final fallback
         logger.warning("All SAM2 methods failed - using simple threshold")
@@ -263,41 +264,47 @@ def _get_mask(self, frame: np.ndarray) -> np.ndarray:
         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
+        """Maximum quality mask cleaning and refinement - CURRENTLY DISABLED."""
+        # THIS METHOD IS DISABLED TO PREVENT DOUBLE CLEANING
+        # The SAM2Handler already applies its own cleaning
+        logger.info("Maximum quality mask cleaning DISABLED to prevent double processing")
+        return mask
+        
+        # ORIGINAL CODE COMMENTED OUT:
+        # 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]:
@@ -466,8 +473,8 @@ def _prog(p, d):
                 if not ok:
                     break
 
-                # --- ENHANCED SAM2 segmentation with full-body detection ---
-                logger.info(f"Processing frame {frame_idx} with enhanced SAM2")
+                # --- SAM2 segmentation WITHOUT double cleaning ---
+                logger.info(f"Processing frame {frame_idx} with SAM2 (no double cleaning)")
                 mask = self._get_mask(frame)
 
                 # --- MatAnyone bootstrap exactly once (first frame) ---