Update processing/two_stage/quality_manager.py

processing/two_stage/quality_manager.py

@@ -1,6 +1,6 @@
 """
-Enhanced Quality Profile Manager for Two-Stage Processing
-THIS IS A NEW FILE - Create this to separate quality management logic
+Enhanced Quality Profile Manager for Two-Stage Processing - MAXIMUM QUALITY VERSION
+Updated to align with maximum quality SAM2 settings and aggressive mask cleaning
 """
 import os
 import logging
@@ -11,27 +11,27 @@
 
 @dataclass
 class QualityProfile:
-    """Quality profile with all tunable parameters"""
+    """Quality profile with all tunable parameters - optimized for maximum quality"""
     # MatAnyone refinement
     refine_enabled: bool = True
     refine_cadence: int = 3  # Process every N frames (1=every frame)
     refine_start_frame: int = 0  # Start refinement from this frame
     
-    # Mask processing
+    # Mask processing - enhanced for maximum quality
     mask_dilate: int = 5  # Morphological dilation radius
     mask_blur: int = 3  # Gaussian blur for feathering
     mask_threshold: float = 0.5  # Binary threshold
     
-    # Chroma keying
+    # Chroma keying - tightened for better precision
     chroma_tolerance: float = 30.0  # Color distance tolerance
     chroma_softness: float = 10.0  # Edge softness
     spill_suppression: float = 0.3  # Green/blue spill removal strength
     
-    # Hybrid mixing
+    # Hybrid mixing - favor segmentation for better quality
     hybrid_mix: float = 0.7  # 0=pure chroma, 1=pure segmentation
     temporal_smooth: float = 0.2  # Temporal consistency weight
     
-    # Edge refinement
+    # Edge refinement - enhanced for smoother results
     edge_feather: int = 5  # Edge feathering radius
     edge_contrast: float = 1.2  # Edge contrast enhancement
     
@@ -42,19 +42,19 @@ class QualityProfile:
     max_resolution: int = 1080  # Max processing resolution
     batch_size: int = 1  # Frames to process in parallel
 
-# Define quality presets with VISIBLE differentiation
+# Define quality presets with MAXIMUM QUALITY optimizations
 QUALITY_PRESETS = {
     "speed": QualityProfile(
         refine_enabled=True,
         refine_cadence=10,  # Refine every 10th frame only
         refine_start_frame=0,
-        mask_dilate=3,
+        mask_dilate=2,
         mask_blur=1,
         mask_threshold=0.5,
-        chroma_tolerance=40.0,  # Less precise
+        chroma_tolerance=50.0,  # Less precise for speed
         chroma_softness=15.0,
         spill_suppression=0.1,  # Minimal spill work
-        hybrid_mix=0.5,  # Balanced mix
+        hybrid_mix=0.4,  # Favor chroma for speed
         temporal_smooth=0.1,
         edge_feather=2,
         edge_contrast=1.0,
@@ -69,9 +69,9 @@ class QualityProfile:
         mask_dilate=5,
         mask_blur=3,
         mask_threshold=0.5,
-        chroma_tolerance=30.0,
-        chroma_softness=10.0,
-        spill_suppression=0.3,
+        chroma_tolerance=35.0,
+        chroma_softness=8.0,
+        spill_suppression=0.35,
         hybrid_mix=0.7,  # Favor segmentation
         temporal_smooth=0.2,
         edge_feather=5,
@@ -82,41 +82,41 @@ class QualityProfile:
     ),
     "max": QualityProfile(
         refine_enabled=True,
-        refine_cadence=1,  # Refine EVERY frame
+        refine_cadence=1,  # Refine EVERY frame for maximum quality
         refine_start_frame=0,
-        mask_dilate=7,  # Stronger hole filling
-        mask_blur=5,  # Smoother edges
-        mask_threshold=0.45,
-        chroma_tolerance=20.0,  # More precise
-        chroma_softness=5.0,  # Sharper edges
-        spill_suppression=0.6,  # Strong spill removal
-        hybrid_mix=0.85,  # Heavy favor to segmentation
-        temporal_smooth=0.3,  # More temporal consistency
-        edge_feather=7,
-        edge_contrast=1.5,  # Stronger edge enhancement
-        bg_blur_sigma=2.0,  # More background blur
-        max_resolution=1920,  # Full HD+
-        batch_size=1
+        mask_dilate=8,  # MAXIMUM: Aggressive hole filling to eliminate gaps
+        mask_blur=6,  # MAXIMUM: Heavy smoothing for seamless edges
+        mask_threshold=0.4,  # MAXIMUM: Lower threshold to keep more detail
+        chroma_tolerance=25.0,  # MAXIMUM: Tighter tolerance for precision
+        chroma_softness=3.0,  # MAXIMUM: Minimal softness for crisp edges
+        spill_suppression=0.6,  # MAXIMUM: Strong spill removal
+        hybrid_mix=0.9,  # MAXIMUM: Heavy favor to segmentation over chroma
+        temporal_smooth=0.4,  # MAXIMUM: Strong temporal consistency
+        edge_feather=8,  # MAXIMUM: Heavy edge feathering
+        edge_contrast=1.8,  # MAXIMUM: Strong edge enhancement
+        bg_blur_sigma=2.5,  # MAXIMUM: More background blur for seamless blend
+        max_resolution=1920,  # MAXIMUM: Full HD+ processing
+        batch_size=1  # Single frame processing for stability
     )
 }
 
 class QualityManager:
-    """Manages quality profiles and settings"""
+    """Manages quality profiles and settings with maximum quality optimizations"""
     
     def __init__(self):
         self.current_profile: Optional[QualityProfile] = None
-        self.profile_name: str = "balanced"
+        self.profile_name: str = "max"  # Default to maximum quality
         
     def load_profile(self, profile_name: Optional[str] = None) -> QualityProfile:
         """Load quality profile from environment or parameter"""
-        # Check environment variable first
+        # Check environment variable first, default to max quality
         if profile_name is None:
-            profile_name = os.environ.get("BFX_QUALITY", "balanced").lower()
+            profile_name = os.environ.get("BFX_QUALITY", "max").lower()
         
         # Validate and load profile
         if profile_name not in QUALITY_PRESETS:
-            logger.warning(f"Unknown quality profile '{profile_name}', using 'balanced'")
-            profile_name = "balanced"
+            logger.warning(f"Unknown quality profile '{profile_name}', using 'max' for best results")
+            profile_name = "max"
         
         self.profile_name = profile_name
         self.current_profile = QUALITY_PRESETS[profile_name]
@@ -134,6 +134,8 @@ def _profile_summary(self) -> str:
         p = self.current_profile
         return (
             f"refine_every={p.refine_cadence} frames, "
+            f"mask_dilate={p.mask_dilate}, "
+            f"mask_blur={p.mask_blur}, "
             f"spill={p.spill_suppression:.1f}, "
             f"hybrid_mix={p.hybrid_mix:.1f}, "
             f"edge_feather={p.edge_feather}, "
@@ -152,7 +154,7 @@ def should_refine_frame(self, frame_idx: int) -> bool:
         if frame_idx == 0:
             return True
         
-        # Check cadence
+        # Check cadence - for max quality, refine every frame
         return (frame_idx % self.current_profile.refine_cadence) == 0
     
     def get_params(self) -> Dict[str, Any]:
@@ -173,8 +175,36 @@ def log_quality_impact(self, frame_idx: int, processing_time: float):
                 f"Frame {frame_idx}: Skipped refinement (quality={self.profile_name}, "
                 f"time={processing_time:.3f}s)"
             )
+    
+    def get_max_quality_params(self) -> Dict[str, Any]:
+        """Get parameters specifically optimized for maximum quality gap elimination"""
+        max_profile = QUALITY_PRESETS["max"]
+        params = asdict(max_profile)
+        
+        # Add specific maximum quality enhancements
+        params.update({
+            # Aggressive morphological operations
+            "morph_close_kernel": 9,  # Large closing kernel for gap filling
+            "morph_connect_kernel": 7,  # Connect nearby regions
+            "gaussian_heavy_sigma": 2.0,  # Heavy boundary smoothing
+            "gaussian_final_sigma": 1.0,  # Final edge smoothing
+            "dilate_final_kernel": 3,  # Final dilation for coverage
+            
+            # Enhanced SAM2 integration
+            "sam2_points_per_side": 64,  # Double density sampling
+            "sam2_confidence_threshold": 0.05,  # Accept almost everything
+            "sam2_bbox_padding": 0.3,  # 30% padding for full coverage
+            "sam2_refine_iterations": 5,  # Maximum refinement iterations
+            
+            # Quality flags
+            "enable_aggressive_cleaning": True,
+            "enable_heavy_smoothing": True,
+            "enable_gap_elimination": True,
+        })
+        
+        return params
 
-# Global instance
+# Global instance with maximum quality default
 quality_manager = QualityManager()
 
 def get_quality_profile(profile_name: Optional[str] = None) -> QualityProfile:
@@ -183,4 +213,15 @@ def get_quality_profile(profile_name: Optional[str] = None) -> QualityProfile:
 
 def should_refine_frame(frame_idx: int) -> bool:
     """Check if frame should be refined (convenience function)"""
-    return quality_manager.should_refine_frame(frame_idx)
+    return quality_manager.should_refine_frame(frame_idx)
+
+def get_max_quality_settings() -> Dict[str, Any]:
+    """Get maximum quality settings for gap elimination"""
+    return quality_manager.get_max_quality_params()
+
+def is_max_quality_enabled() -> bool:
+    """Check if maximum quality mode is active"""
+    return quality_manager.profile_name == "max"
+
+# Initialize with maximum quality by default
+quality_manager.load_profile("max")