src/models/ensemble.py

# backend/src/models/ensemble.py
"""
Real YOLO Ensemble Detector for Road Damage Detection

Uses the trained yolo11base.pt model for actual inference.
Architecture supports future expansion to 3 models (YOLO11, YOLO12, SAHI)
deployed on HuggingFace Spaces.

Current Configuration:
- yolo11: 1.0 (active - local model)
- yolo12: 0.0 (pending - training)
- sahi: 0.0 (pending - training)
"""

import os
import numpy as np
from typing import List, Dict, Any, Optional
import logging

logger = logging.getLogger(__name__)

# Configure logging to show in console
logging.basicConfig(level=logging.INFO)

# Try to import YOLO from ultralytics
try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
    print("[ENSEMBLE] ultralytics YOLO imported successfully")
except ImportError:
    YOLO_AVAILABLE = False
    print("[ENSEMBLE] WARNING: ultralytics not installed, using mock detector")
    logger.warning("[WARN] ultralytics not installed, using mock detector")


class EnsembleDetector:
    """
    Real YOLO-based ensemble detector for road damage.
    
    Uses local YOLO11 model for inference. When additional models
    are trained and deployed to HuggingFace Spaces, they can be
    added to the ensemble with appropriate weights.
    """
    
    def __init__(self, model_path: str = None):
        """
        Initialize the ensemble detector.
        
        Args:
            model_path: Path to YOLO model weights. 
                       Defaults to models/yolo11base.pt
        """
        self.model = None
        self.model_loaded = False
        
        # Model weights configuration
        self.weights = {
            "yolo11": 1.0,   # Active - local model
            "yolo12": 0.0,   # Pending
            "sahi": 0.0      # Pending
        }
        
        # Class names from RDD2022 dataset
        self.class_names = {
            0: "D00",  # Longitudinal Crack
            1: "D10",  # Transverse Crack
            2: "D20",  # Alligator Crack
            3: "D40",  # Pothole
        }
        
        self.class_display_names = {
            "D00": "Longitudinal Crack",
            "D10": "Transverse Crack",
            "D20": "Alligator Crack",
            "D40": "Pothole",
        }
        
        # Determine model path
        if model_path is None:
            # Try multiple possible paths (local dev + HuggingFace Docker)
            possible_paths = [
                # HuggingFace Docker container paths (WORKDIR=/app)
                "/app/models/yolo11base.pt",
                "models/yolo11base.pt",
                # Local development paths
                "../models/yolo11base.pt",
                os.path.join(os.path.dirname(__file__), "../../models/yolo11base.pt"),
                os.path.join(os.path.dirname(__file__), "../../../models/yolo11base.pt"),
                "D:/Gitrepo/road-damage/models/yolo11base.pt",
            ]
            print(f"[MODEL] Searching for model in paths: {possible_paths}")
            for path in possible_paths:
                if os.path.exists(path):
                    model_path = path
                    print(f"[MODEL] Found model at: {path}")
                    break
            if model_path is None:
                print(f"[MODEL] WARNING: Model not found in any path!")
        
        # Load model
        if YOLO_AVAILABLE and model_path and os.path.exists(model_path):
            try:
                logger.info(f"[INFO] Loading YOLO model from: {model_path}")
                self.model = YOLO(model_path)
                self.model_loaded = True
                logger.info("[OK] YOLO model loaded successfully!")
                
                # Get class names from model if available
                if hasattr(self.model, 'names') and self.model.names:
                    self.class_names = self.model.names
                    logger.info(f"[INFO] Model classes: {self.class_names}")
                    
            except Exception as e:
                logger.error(f"[ERROR] Failed to load YOLO model: {e}")
                self.model_loaded = False
        else:
            if not YOLO_AVAILABLE:
                logger.warning("[WARN] ultralytics not available")
            elif not model_path:
                logger.warning("[WARN] No model path specified")
            else:
                logger.warning(f"[WARN] Model file not found: {model_path}")
    
    def predict(self, image: np.ndarray, conf_threshold: float = 0.25) -> List[Dict]:
        """
        Run inference on an image.
        
        Args:
            image: numpy array (HWC format, RGB or BGR)
            conf_threshold: Minimum confidence threshold
            
        Returns:
            List of detection dictionaries with keys:
            - box: [x1, y1, x2, y2]
            - class_name: str
            - class_id: int
            - confidence: float
            - votes: int (number of models that detected this)
        """
        
        if not self.model_loaded:
            logger.warning("[WARN] Model not loaded, returning empty results")
            return []
        
        try:
            # Run inference
            results = self.model(
                image,
                conf=conf_threshold,
                verbose=False
            )
            
            detections = []
            
            for result in results:
                boxes = result.boxes
                
                if boxes is None or len(boxes) == 0:
                    continue
                
                for i in range(len(boxes)):
                    # Get bounding box
                    box = boxes.xyxy[i].cpu().numpy()
                    x1, y1, x2, y2 = box
                    
                    # Get confidence
                    conf = float(boxes.conf[i].cpu().numpy())
                    
                    # Get class
                    cls_id = int(boxes.cls[i].cpu().numpy())
                    
                    # Get class name
                    if cls_id in self.class_names:
                        cls_name = self.class_names[cls_id]
                    else:
                        cls_name = f"class_{cls_id}"
                    
                    # Get display name
                    display_name = self.class_display_names.get(cls_name, cls_name)
                    
                    detection = {
                        "box": [float(x1), float(y1), float(x2), float(y2)],
                        "class_name": display_name,
                        "class_code": cls_name,
                        "class_id": cls_id,
                        "confidence": conf,
                        "votes": 1,  # Single model for now
                        "model": "yolo11"
                    }
                    
                    detections.append(detection)
            
            logger.info(f"[DETECT] YOLO11 detected {len(detections)} objects")
            return detections
            
        except Exception as e:
            logger.error(f"[ERROR] Inference failed: {e}")
            import traceback
            traceback.print_exc()
            return []
    
    def predict_with_ensemble(
        self, 
        image: np.ndarray,
        conf_threshold: float = 0.25
    ) -> List[Dict]:
        """
        Run ensemble inference (future: multiple models).
        
        Currently only uses YOLO11 since other models are pending.
        When YOLO12 and SAHI are ready, this method will:
        1. Call all 3 models in parallel
        2. Merge overlapping detections (NMS)
        3. Apply weighted voting
        4. Return only detections agreed by 2+ models
        
        Args:
            image: numpy array
            conf_threshold: Minimum confidence
            
        Returns:
            List of ensemble-merged detections
        """
        
        # For now, just use YOLO11
        yolo11_results = self.predict(image, conf_threshold)
        
        # Future: Add YOLO12 and SAHI results
        # yolo12_results = self._call_hf_space("yolo12", image)
        # sahi_results = self._call_hf_space("sahi", image)
        
        # Future: Merge with NMS and voting
        # merged = self._merge_detections([yolo11_results, yolo12_results, sahi_results])
        
        return yolo11_results
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get information about loaded models."""
        return {
            "yolo11": {
                "loaded": self.model_loaded,
                "weight": self.weights["yolo11"],
                "type": "local"
            },
            "yolo12": {
                "loaded": False,
                "weight": self.weights["yolo12"],
                "type": "hf_space",
                "status": "pending_training"
            },
            "sahi": {
                "loaded": False,
                "weight": self.weights["sahi"],
                "type": "hf_space", 
                "status": "pending_training"
            }
        }


class SeverityClassifier:
    """Classify damage severity based on area and confidence."""
    
    def __init__(self):
        # Thresholds for severity classification
        self.thresholds = {
            "light": 0.05,    # < 5% of image
            "medium": 0.15,   # 5-15% of image
            # > 15% = heavy
        }
    
    def classify(self, area_ratio: float, confidence: float) -> str:
        """
        Classify severity based on damage area relative to image.
        
        Args:
            area_ratio: Damage area / image area
            confidence: Detection confidence (0-1)
            
        Returns:
            Severity string: "light", "medium", or "heavy"
        """
        # Adjust thresholds based on confidence
        # Lower confidence = more conservative severity
        conf_factor = min(1.0, confidence / 0.5)
        
        effective_ratio = area_ratio * conf_factor
        
        if effective_ratio > self.thresholds["medium"]:
            return "heavy"
        elif effective_ratio > self.thresholds["light"]:
            return "medium"
        else:
            return "light"


class TTAProcessor:
    """Test-Time Augmentation processor (optional enhancement)."""
    
    def __init__(self):
        self.enabled = False  # Disabled by default for speed
    
    def predict(self, image: np.ndarray) -> List[Dict]:
        """
        Run TTA inference (currently disabled).
        
        When enabled, this applies augmentations and averages results
        for more robust detections at the cost of speed.
        """
        if not self.enabled:
            return []
        
        # Future: Implement TTA
        # - Horizontal flip
        # - Multi-scale
        # - Merge results
        return []