Create processing/two_stage/quality_profiles.py

processing/two_stage/quality_profiles.py

@@ -0,0 +1,186 @@
+"""
+Enhanced Quality Profile Manager for Two-Stage Processing
+THIS IS A NEW FILE - Create this to separate quality management logic
+"""
+import os
+import logging
+from typing import Dict, Any, Optional
+from dataclasses import dataclass, asdict
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class QualityProfile:
+    """Quality profile with all tunable parameters"""
+    # 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_dilate: int = 5  # Morphological dilation radius
+    mask_blur: int = 3  # Gaussian blur for feathering
+    mask_threshold: float = 0.5  # Binary threshold
+    
+    # Chroma keying
+    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_mix: float = 0.7  # 0=pure chroma, 1=pure segmentation
+    temporal_smooth: float = 0.2  # Temporal consistency weight
+    
+    # Edge refinement
+    edge_feather: int = 5  # Edge feathering radius
+    edge_contrast: float = 1.2  # Edge contrast enhancement
+    
+    # Background
+    bg_blur_sigma: float = 1.0  # Background blur strength
+    
+    # Performance
+    max_resolution: int = 1080  # Max processing resolution
+    batch_size: int = 1  # Frames to process in parallel
+
+# Define quality presets with VISIBLE differentiation
+QUALITY_PRESETS = {
+    "speed": QualityProfile(
+        refine_enabled=True,
+        refine_cadence=10,  # Refine every 10th frame only
+        refine_start_frame=0,
+        mask_dilate=3,
+        mask_blur=1,
+        mask_threshold=0.5,
+        chroma_tolerance=40.0,  # Less precise
+        chroma_softness=15.0,
+        spill_suppression=0.1,  # Minimal spill work
+        hybrid_mix=0.5,  # Balanced mix
+        temporal_smooth=0.1,
+        edge_feather=2,
+        edge_contrast=1.0,
+        bg_blur_sigma=0.5,
+        max_resolution=720,
+        batch_size=2
+    ),
+    "balanced": QualityProfile(
+        refine_enabled=True,
+        refine_cadence=3,  # Refine every 3rd frame
+        refine_start_frame=0,
+        mask_dilate=5,
+        mask_blur=3,
+        mask_threshold=0.5,
+        chroma_tolerance=30.0,
+        chroma_softness=10.0,
+        spill_suppression=0.3,
+        hybrid_mix=0.7,  # Favor segmentation
+        temporal_smooth=0.2,
+        edge_feather=5,
+        edge_contrast=1.2,
+        bg_blur_sigma=1.0,
+        max_resolution=1080,
+        batch_size=1
+    ),
+    "max": QualityProfile(
+        refine_enabled=True,
+        refine_cadence=1,  # Refine EVERY frame
+        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
+    )
+}
+
+class QualityManager:
+    """Manages quality profiles and settings"""
+    
+    def __init__(self):
+        self.current_profile: Optional[QualityProfile] = None
+        self.profile_name: str = "balanced"
+        
+    def load_profile(self, profile_name: Optional[str] = None) -> QualityProfile:
+        """Load quality profile from environment or parameter"""
+        # Check environment variable first
+        if profile_name is None:
+            profile_name = os.environ.get("BFX_QUALITY", "balanced").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"
+        
+        self.profile_name = profile_name
+        self.current_profile = QUALITY_PRESETS[profile_name]
+        
+        # Log the switch with details
+        logger.info(f"Quality switched to '{profile_name}' ⇒ {self._profile_summary()}")
+        
+        return self.current_profile
+    
+    def _profile_summary(self) -> str:
+        """Generate a summary of key profile parameters"""
+        if not self.current_profile:
+            return "No profile loaded"
+        
+        p = self.current_profile
+        return (
+            f"refine_every={p.refine_cadence} frames, "
+            f"spill={p.spill_suppression:.1f}, "
+            f"hybrid_mix={p.hybrid_mix:.1f}, "
+            f"edge_feather={p.edge_feather}, "
+            f"bg_blur={p.bg_blur_sigma:.1f}"
+        )
+    
+    def should_refine_frame(self, frame_idx: int) -> bool:
+        """Check if current frame should be refined with MatAnyone"""
+        if not self.current_profile or not self.current_profile.refine_enabled:
+            return False
+        
+        if frame_idx < self.current_profile.refine_start_frame:
+            return False
+        
+        # Always refine first frame for bootstrap
+        if frame_idx == 0:
+            return True
+        
+        # Check cadence
+        return (frame_idx % self.current_profile.refine_cadence) == 0
+    
+    def get_params(self) -> Dict[str, Any]:
+        """Get current profile as dictionary"""
+        if not self.current_profile:
+            self.load_profile()
+        return asdict(self.current_profile)
+    
+    def log_quality_impact(self, frame_idx: int, processing_time: float):
+        """Log quality impact for monitoring"""
+        if self.should_refine_frame(frame_idx):
+            logger.debug(
+                f"Frame {frame_idx}: Refined (quality={self.profile_name}, "
+                f"time={processing_time:.3f}s)"
+            )
+        else:
+            logger.debug(
+                f"Frame {frame_idx}: Skipped refinement (quality={self.profile_name}, "
+                f"time={processing_time:.3f}s)"
+            )
+
+# Global instance
+quality_manager = QualityManager()
+
+def get_quality_profile(profile_name: Optional[str] = None) -> QualityProfile:
+    """Get quality profile (convenience function)"""
+    return quality_manager.load_profile(profile_name)
+
+def should_refine_frame(frame_idx: int) -> bool:
+    """Check if frame should be refined (convenience function)"""
+    return quality_manager.should_refine_frame(frame_idx)