model_loader.py

#!/usr/bin/env python3
"""
Robust Model Loader for PPE Detection
Handles model loading with fallbacks and error recovery
"""

import os
import torch
import logging
from pathlib import Path
from ultralytics import YOLO
import requests
import zipfile

logger = logging.getLogger(__name__)

class RobustModelLoader:
    def __init__(self):
        self.model_paths = [
            # Improved models (highest priority)
            'best.pt',
            'model/PPE_DETECTION/runs/detect/*/weights/best.pt',
            
            # Existing models
            'weights/best.pt',
            'model/PPE_DETECTION/Models/best_new.pt',
            'model/PPE_DETECTION/Models/best.pt',
            'best.pt',
            
            # Fallback models
            'yolov8n.pt',
            'yolov8s.pt'
        ]
        
        self.backup_urls = {
            'yolov8n.pt': 'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt',
            'yolov8s.pt': 'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt'
        }
    
    def find_best_model(self):
        """Find the best available model"""
        logger.info("🔍 Searching for best available model...")
        
        for path_pattern in self.model_paths:
            if '*' in path_pattern:
                # Handle glob patterns
                from glob import glob
                matches = glob(path_pattern)
                if matches:
                    # Use most recent model
                    best_match = max(matches, key=os.path.getmtime)
                    if self.validate_model(best_match):
                        logger.info(f"✅ Found improved model: {best_match}")
                        return best_match
            else:
                if os.path.exists(path_pattern):
                    if self.validate_model(path_pattern):
                        logger.info(f"✅ Found model: {path_pattern}")
                        return path_pattern
        
        # No model found, try to download fallback
        logger.warning("⚠️ No local model found, attempting to download fallback...")
        return self.download_fallback_model()
    
    def validate_model(self, model_path):
        """Validate that a model file is usable"""
        try:
            if not os.path.exists(model_path):
                return False
            
            # Check file size (should be > 1MB for a real model)
            if os.path.getsize(model_path) < 1024 * 1024:
                logger.warning(f"⚠️ Model file too small: {model_path}")
                return False
            
            # Try to load the model
            test_model = YOLO(model_path)
            
            # Quick test inference
            import numpy as np
            dummy_img = np.zeros((640, 640, 3), dtype=np.uint8)
            results = test_model(dummy_img, verbose=False)
            
            logger.info(f"✅ Model validation successful: {model_path}")
            return True
            
        except Exception as e:
            logger.error(f"❌ Model validation failed for {model_path}: {e}")
            return False
    
    def download_fallback_model(self):
        """Download a fallback model if none available"""
        for model_name, url in self.backup_urls.items():
            try:
                logger.info(f"📥 Downloading fallback model: {model_name}")
                
                response = requests.get(url, stream=True)
                response.raise_for_status()
                
                with open(model_name, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)
                
                if self.validate_model(model_name):
                    logger.info(f"✅ Successfully downloaded: {model_name}")
                    return model_name
                    
            except Exception as e:
                logger.error(f"❌ Failed to download {model_name}: {e}")
                continue
        
        logger.error("❌ Could not find or download any usable model!")
        return None
    
    def load_model_with_fallback(self):
        """Load model with comprehensive error handling"""
        model_path = self.find_best_model()
        
        if not model_path:
            raise RuntimeError("No usable model found!")
        
        try:
            logger.info(f"🤖 Loading model: {model_path}")
            model = YOLO(model_path)
            
            # Configure device
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            logger.info(f"🔧 Using device: {device}")
            
            try:
                model.to(device)
                if device == 'cuda' and hasattr(model.model, 'half'):
                    model.model.half()  # Use FP16 for speed
            except Exception as e:
                logger.warning(f"⚠️ Could not optimize model: {e}")
            
            # Test the model
            import numpy as np
            test_img = np.zeros((640, 640, 3), dtype=np.uint8)
            results = model(test_img, verbose=False)
            
            logger.info("✅ Model loaded and tested successfully!")
            
            # Print model info
            if 'advanced_training' in model_path:
                logger.info("🚀 Using IMPROVED AI model!")
            elif 'yolov8' in model_path:
                logger.info("⚠️ Using basic YOLO model - consider training a custom PPE model")
            else:
                logger.info("📊 Using custom PPE model")
            
            return model, model_path
            
        except Exception as e:
            logger.error(f"❌ Failed to load model {model_path}: {e}")
            
            # Try fallback
            if 'yolov8' not in model_path:
                logger.info("🔄 Trying fallback model...")
                fallback_path = self.download_fallback_model()
                if fallback_path:
                    return self.load_model_with_fallback()
            
            raise RuntimeError(f"Could not load any model: {e}")

def load_ppe_model():
    """Main function to load PPE detection model"""
    loader = RobustModelLoader()
    return loader.load_model_with_fallback()