core/models.py

"""
Core Models Module - Simplified redirect to working model loader
"""

import logging

logger = logging.getLogger(__name__)

class ModelManager:
    """
    Compatibility wrapper that redirects to the working ModelLoader
    Provides the same interface that CoreVideoProcessor expects
    """
    
    def __init__(self):
        self._loader = None
        self._device_mgr = None
        self._memory_mgr = None
        
    def _get_loader(self):
        """Lazy initialization of model loader"""
        if self._loader is None:
            try:
                from models.loaders.model_loader import ModelLoader
                from utils.hardware.device_manager import DeviceManager
                from utils.system.memory_manager import MemoryManager
                
                self._device_mgr = DeviceManager()
                self._memory_mgr = MemoryManager()
                self._loader = ModelLoader(self._device_mgr, self._memory_mgr)
                
            except ImportError as e:
                logger.error(f"Failed to import ModelLoader dependencies: {e}")
                # Create a dummy loader that returns None for everything
                class DummyLoader:
                    def get_sam2(self): return None
                    def get_matanyone(self): return None
                    def load_all_models(self, *args, **kwargs): return None, None
                    def cleanup(self): pass
                    
                self._loader = DummyLoader()
                
        return self._loader
    
    def load_all(self):
        """Load all models (SAM2 and MatAnyone)"""
        try:
            loader = self._get_loader()
            sam2, matanyone = loader.load_all_models()
            logger.info(f"Models loaded - SAM2: {sam2 is not None}, MatAnyone: {matanyone is not None}")
            return sam2, matanyone
        except Exception as e:
            logger.error(f"Failed to load models: {e}")
            return None, None
    
    def get_sam2(self):
        """Get SAM2 model/predictor"""
        loader = self._get_loader()
        return loader.get_sam2() if hasattr(loader, 'get_sam2') else None
    
    def get_matanyone(self):
        """Get MatAnyone model"""
        loader = self._get_loader()
        return loader.get_matanyone() if hasattr(loader, 'get_matanyone') else None
    
    def cleanup(self):
        """Cleanup models and free memory"""
        loader = self._get_loader()
        if hasattr(loader, 'cleanup'):
            loader.cleanup()
    
    def process_frame(self, image, mask=None):
        """Process a single frame - for compatibility"""
        # This method was in the old core/models.py
        # We'll just return a dummy result since the actual processing
        # happens in CoreVideoProcessor
        logger.warning("ModelManager.process_frame called - this is deprecated")
        return {
            "alpha": mask if mask is not None else image[:,:,0],
            "foreground": image,
            "confidence": [0.5]
        }

# For backward compatibility - some code might import these directly
ModelType = None  # Not needed with new system
ModelFactory = ModelManager  # Alias for compatibility
ModelConfig = None  # Configuration now handled in model_loader.py

__all__ = ['ModelManager']