model/model_loader.py

"""
Model loading module for food detection.
Loads local model file provided by user.
Supports Ultralytics YOLO (.pt) and TensorFlow (.h5) models.
"""

import logging
import os
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


class ModelLoader:
    """Singleton loader for local model files."""
    _instance = None
    _model = None
    _model_path = None
    _conf_threshold = 0.25
    _iou_threshold = 0.5
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def load_model(self, model_path: Optional[str] = None, conf: float = 0.25, iou: float = 0.5):
        """
        Load model from local file.
        
        Args:
            model_path: Path to model file (.pt for YOLO, .h5 for TensorFlow)
            conf: Confidence threshold
            iou: IOU threshold for NMS
            
        Returns:
            Loaded model
            
        Raises:
            FileNotFoundError: If model file not found
            ValueError: If unsupported model format
        """
        if model_path is None:
            # Try common locations and env variable
            common_paths = [
                os.getenv('MODEL_PATH'),
                'models/weight5_multi_yolov11L.pt',
                'models/best.pt',
                'models/weights4_fridge_vision_yolov8l.pt',
                'models/weights3_fridge_vision_yolov8l.pt',
                'models/model.pt',
                'model.pt',
            ]
            model_path = next((p for p in common_paths if p and os.path.exists(p)), None)
            
            if not model_path:
                raise FileNotFoundError(
                    "❌ Model file not found. Please:\n"
                    "1. Set MODEL_PATH environment variable, OR\n"
                    "2. Place model in models/ directory (weight5_multi_yolov11L.pt, best.pt, or model.pt)"
                )
        
        # Return cached model if same path
        if self._model is not None and self._model_path == model_path:
            return self._model
        
        model_path = str(model_path)
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model file not found: {model_path}")
        
        logger.info(f"Loading model from: {model_path}")
        
        try:
            model_ext = Path(model_path).suffix.lower()
            self._conf_threshold = conf
            self._iou_threshold = iou
            
            if model_ext == '.pt':
                # Ultralytics YOLO model (YOLOv8/YOLOv11)
                from ultralytics import YOLO

                self._model = YOLO(model_path)
                if hasattr(self._model, 'overrides') and isinstance(self._model.overrides, dict):
                    self._model.overrides['conf'] = conf
                    self._model.overrides['iou'] = iou
                logger.info(f"✅ Ultralytics YOLO model loaded (conf={conf}, iou={iou})")
                
            elif model_ext == '.h5':
                # TensorFlow/Keras model
                try:
                    import tensorflow as tf
                    self._model = tf.keras.models.load_model(model_path)
                    logger.info("✅ TensorFlow model loaded")
                except ImportError:
                    raise ImportError("TensorFlow required for .h5 models")
                    
            else:
                raise ValueError(f"Unsupported model format: {model_ext}. Use .pt or .h5")
            
            self._model_path = model_path
            return self._model
            
        except Exception as e:
            logger.error(f"❌ Failed to load model: {e}")
            self._model = None
            self._model_path = None
            raise
    
    def get_model(self, model_path: Optional[str] = None):
        """Get loaded model or load it if not already loaded."""
        if self._model is None:
            return self.load_model(model_path)
        if model_path is not None and str(model_path) != self._model_path:
            return self.load_model(model_path, conf=self._conf_threshold, iou=self._iou_threshold)
        return self._model

    @property
    def model_path(self) -> Optional[str]:
        """Return currently loaded model path."""
        return self._model_path

    @property
    def conf_threshold(self) -> float:
        """Return confidence threshold stored in loader."""
        return self._conf_threshold

    @property
    def iou_threshold(self) -> float:
        """Return IOU threshold stored in loader."""
        return self._iou_threshold

    @property
    def device(self) -> str:
        """Best-effort device info for compatibility with old callers."""
        if self._model is None:
            return "unknown"
        model_obj = getattr(self._model, 'model', None)
        if model_obj is not None and hasattr(model_obj, 'device'):
            return str(model_obj.device)
        return "cpu"

    def set_confidence_threshold(self, conf: float):
        """Set default confidence threshold for compatibility with old callers."""
        self._conf_threshold = conf
        if self._model is not None and hasattr(self._model, 'overrides') and isinstance(self._model.overrides, dict):
            self._model.overrides['conf'] = conf

    def set_iou_threshold(self, iou: float):
        """Set default IOU threshold for compatibility with old callers."""
        self._iou_threshold = iou
        if self._model is not None and hasattr(self._model, 'overrides') and isinstance(self._model.overrides, dict):
            self._model.overrides['iou'] = iou
    
    def clear_cache(self):
        """Clear cached model (use when switching models)."""
        self._model = None
        self._model_path = None
        logger.info("Model cache cleared")


def get_model_loader(
    model_path: Optional[str] = None,
    conf: float = 0.25,
    iou: float = 0.5,
) -> ModelLoader:
    """Get singleton ModelLoader instance, optionally preloading a model."""
    loader = ModelLoader()
    if model_path is not None and (loader.model_path != str(model_path) or loader._model is None):
        loader.load_model(model_path=model_path, conf=conf, iou=iou)
    return loader