Update config/app_config.py

config/app_config.py

@@ -46,6 +46,34 @@ class ProcessingConfig:
     matanyone_precision: str = os.getenv('MATANYONE_PRECISION', 'fp32')  # fp16, fp32
     model_device: str = os.getenv('MODEL_DEVICE', 'auto')  # auto, cuda, cpu
     
+    # MatAnyone specific settings
+    matanyone_enabled: bool = os.getenv('MATANYONE_ENABLED', 'false').lower() == 'true'
+    matanyone_model_path: str = os.getenv('MATANYONE_MODEL_PATH', 'models/matanyone/checkpoint.pth')
+    matanyone_threshold: float = float(os.getenv('MATANYONE_THRESHOLD', '0.5'))
+    matanyone_hair_refinement: bool = os.getenv('MATANYONE_HAIR_REFINEMENT', 'true').lower() == 'true'
+    matanyone_edge_enhancement: bool = os.getenv('MATANYONE_EDGE_ENHANCEMENT', 'true').lower() == 'true'
+    matanyone_detail_refinement: bool = os.getenv('MATANYONE_DETAIL_REFINEMENT', 'true').lower() == 'true'
+    matanyone_morphology_ops: bool = os.getenv('MATANYONE_MORPHOLOGY_OPS', 'true').lower() == 'true'
+    matanyone_morphology_kernel_size: int = int(os.getenv('MATANYONE_MORPHOLOGY_KERNEL_SIZE', '5'))
+    
+    # Component paths for separate mask files
+    hair_masks_dir: str = os.getenv('HAIR_MASKS_DIR', 'assets/hair_masks')
+    edge_masks_dir: str = os.getenv('EDGE_MASKS_DIR', 'assets/edge_masks')
+    detail_masks_dir: str = os.getenv('DETAIL_MASKS_DIR', 'assets/detail_masks')
+    use_component_masks: bool = os.getenv('USE_COMPONENT_MASKS', 'false').lower() == 'true'
+    
+    # MatAnyone component weights
+    matanyone_weight_base: float = float(os.getenv('MATANYONE_WEIGHT_BASE', '1.0'))
+    matanyone_weight_hair: float = float(os.getenv('MATANYONE_WEIGHT_HAIR', '1.2'))
+    matanyone_weight_edge: float = float(os.getenv('MATANYONE_WEIGHT_EDGE', '1.5'))
+    matanyone_weight_detail: float = float(os.getenv('MATANYONE_WEIGHT_DETAIL', '1.1'))
+    
+    # MatAnyone processing modes
+    matanyone_processing_mode: str = os.getenv('MATANYONE_PROCESSING_MODE', 'refine')  # refine, replace, blend
+    matanyone_blend_alpha: float = float(os.getenv('MATANYONE_BLEND_ALPHA', '0.7'))
+    matanyone_trimap_enabled: bool = os.getenv('MATANYONE_TRIMAP_ENABLED', 'false').lower() == 'true'
+    matanyone_trimap_dilation: int = int(os.getenv('MATANYONE_TRIMAP_DILATION', '10'))
+    
     # Processing settings
     temporal_consistency: bool = os.getenv('TEMPORAL_CONSISTENCY', 'true').lower() == 'true'
     edge_refinement: bool = os.getenv('EDGE_REFINEMENT', 'true').lower() == 'true'
@@ -145,6 +173,22 @@ def _validate_config(self):
             logger.warning(f"Invalid precision: {self.matanyone_precision}. Setting to 'fp32'.")
             self.matanyone_precision = 'fp32'
         
+        # Validate MatAnyone settings
+        self.matanyone_threshold = max(0.0, min(1.0, self.matanyone_threshold))
+        self.matanyone_weight_base = max(0.1, min(2.0, self.matanyone_weight_base))
+        self.matanyone_weight_hair = max(0.1, min(2.0, self.matanyone_weight_hair))
+        self.matanyone_weight_edge = max(0.1, min(2.0, self.matanyone_weight_edge))
+        self.matanyone_weight_detail = max(0.1, min(2.0, self.matanyone_weight_detail))
+        self.matanyone_blend_alpha = max(0.0, min(1.0, self.matanyone_blend_alpha))
+        self.matanyone_morphology_kernel_size = max(3, min(15, self.matanyone_morphology_kernel_size))
+        self.matanyone_trimap_dilation = max(1, min(50, self.matanyone_trimap_dilation))
+        
+        # Validate MatAnyone processing mode
+        valid_modes = ['refine', 'replace', 'blend']
+        if self.matanyone_processing_mode not in valid_modes:
+            logger.warning(f"Invalid MatAnyone processing mode: {self.matanyone_processing_mode}. Setting to 'refine'.")
+            self.matanyone_processing_mode = 'refine'
+        
         # Validate output settings
         valid_formats = ['mp4', 'avi', 'mov', 'webm', 'mkv']
         if self.output_format not in valid_formats:
@@ -183,6 +227,14 @@ def _create_directories(self):
             Path(self.output_dir) / 'two_stage'
         ]
         
+        # Add MatAnyone component directories if enabled
+        if self.use_component_masks:
+            directories.extend([
+                self.hair_masks_dir,
+                self.edge_masks_dir,
+                self.detail_masks_dir
+            ])
+        
         for directory in directories:
             try:
                 Path(directory).mkdir(parents=True, exist_ok=True)
@@ -255,6 +307,35 @@ def from_yaml(cls, filepath: str) -> 'ProcessingConfig':
             config_dict = yaml.safe_load(f)
         return cls(**config_dict)
     
+    def get_matanyone_config(self) -> Dict[str, Any]:
+        """Get MatAnyone-specific configuration as a dictionary"""
+        return {
+            'enabled': self.matanyone_enabled,
+            'model_path': self.matanyone_model_path,
+            'threshold': self.matanyone_threshold,
+            'edge_refinement': self.matanyone_edge_enhancement,
+            'hair_refinement': self.matanyone_hair_refinement,
+            'detail_refinement': self.matanyone_detail_refinement,
+            'morphology_ops': self.matanyone_morphology_ops,
+            'morphology_kernel_size': self.matanyone_morphology_kernel_size,
+            'processing_mode': self.matanyone_processing_mode,
+            'blend_alpha': self.matanyone_blend_alpha,
+            'trimap_enabled': self.matanyone_trimap_enabled,
+            'trimap_dilation': self.matanyone_trimap_dilation,
+            'component_weights': {
+                'base': self.matanyone_weight_base,
+                'hair': self.matanyone_weight_hair,
+                'edge': self.matanyone_weight_edge,
+                'detail': self.matanyone_weight_detail
+            },
+            'component_paths': {
+                'hair': self.hair_masks_dir if self.use_component_masks else None,
+                'edge': self.edge_masks_dir if self.use_component_masks else None,
+                'detail': self.detail_masks_dir if self.use_component_masks else None
+            },
+            'precision': self.matanyone_precision
+        }
+    
     def get_quality_settings(self) -> Dict[str, Any]:
         """Get quality-specific settings based on preset"""
         quality_maps = {
@@ -265,7 +346,8 @@ def get_quality_settings(self) -> Dict[str, Any]:
                 'temporal_consistency': False,
                 'model_precision': 'fp16',
                 'batch_size': min(self.batch_size * 2, 16),
-                'output_quality_params': '-preset ultrafast -crf 28'
+                'output_quality_params': '-preset ultrafast -crf 28',
+                'matanyone_enabled': False  # Disable MatAnyone for fast mode
             },
             'balanced': {
                 'keyframe_interval': self.keyframe_interval,
@@ -274,7 +356,8 @@ def get_quality_settings(self) -> Dict[str, Any]:
                 'temporal_consistency': True,
                 'model_precision': 'fp32',
                 'batch_size': self.batch_size,
-                'output_quality_params': '-preset medium -crf 23'
+                'output_quality_params': '-preset medium -crf 23',
+                'matanyone_enabled': self.matanyone_enabled
             },
             'high': {
                 'keyframe_interval': max(self.keyframe_interval // 2, 1),
@@ -283,7 +366,8 @@ def get_quality_settings(self) -> Dict[str, Any]:
                 'temporal_consistency': True,
                 'model_precision': 'fp32',
                 'batch_size': max(self.batch_size // 2, 1),
-                'output_quality_params': '-preset slow -crf 18'
+                'output_quality_params': '-preset slow -crf 18',
+                'matanyone_enabled': True  # Enable MatAnyone for high quality
             },
             'ultra': {
                 'keyframe_interval': 1,
@@ -292,7 +376,10 @@ def get_quality_settings(self) -> Dict[str, Any]:
                 'temporal_consistency': True,
                 'model_precision': 'fp32',
                 'batch_size': 1,
-                'output_quality_params': '-preset veryslow -crf 15'
+                'output_quality_params': '-preset veryslow -crf 15',
+                'matanyone_enabled': True,  # Always enable MatAnyone for ultra
+                'matanyone_morphology_ops': True,
+                'matanyone_trimap_enabled': True
             }
         }
         
@@ -330,7 +417,8 @@ def is_high_performance_mode(self) -> bool:
             self.quality_preset in ['high', 'ultra'] and
             self.edge_refinement and
             self.temporal_consistency and
-            self.keyframe_interval <= 3
+            self.keyframe_interval <= 3 and
+            self.matanyone_enabled  # Include MatAnyone in high performance check
         )
     
     def get_memory_limits(self) -> Dict[str, Any]:
@@ -368,6 +456,13 @@ def validate_for_production(self) -> List[str]:
         if self.max_concurrent_processes > 4:
             warnings.append("High concurrent processes may cause instability")
         
+        # MatAnyone-specific warnings
+        if self.matanyone_enabled and self.quality_preset == 'fast':
+            warnings.append("MatAnyone is enabled with 'fast' quality preset (may impact performance)")
+        
+        if self.use_component_masks and not Path(self.hair_masks_dir).exists():
+            warnings.append(f"Component masks enabled but directory not found: {self.hair_masks_dir}")
+        
         return warnings
 
 # Singleton instance for application-wide use