inference/model_inference.py

"""
Food detection inference pipeline (Updated for YOLOv11L).
Uses local YOLOv11 model (.pt) with explicit preprocessing.
Supports single and batch inference.
"""

import logging
import cv2
import numpy as np
from typing import List, Dict, Tuple, Optional, Union

from ultralytics import YOLO

logger = logging.getLogger(__name__)


class FoodDetectionInference:
    """
    Food detection pipeline using YOLOv11 model (YOLOv11L).
    Matches Roboflow preprocessing: Auto-orient, Resize to 640x640.
    """
    
    # Fridge-specific food classes (30 classes from training)
    DEFAULT_CLASSES = [
        'apple', 'banana', 'blueberry', 'bread', 'brinjal', 'butter', 'cabbage',
        'capsicum', 'carrot', 'cheese', 'chicken', 'chocolate', 'corn', 'cucumber',
        'egg', 'flour', 'fresh cream', 'ginger', 'green beans', 'green chilly',
        'green leaves', 'lemon', 'meat', 'milk', 'mushroom', 'potato', 'shrimp',
        'strawberry', 'sweet potato', 'tomato'
    ]
    
    # Preprocessing constants (matches Roboflow training)
    IMG_SIZE = 640
    CONF_THRESHOLD = 0.25
    IOU_THRESHOLD = 0.7
    
    def __init__(self, model_path: Optional[str] = None, conf: float = 0.25, iou: float = 0.7):
        """
        Initialize detection pipeline with YOLOv11 model.

        Args:
            model_path: Path to .pt model file (e.g., models/weight5_multi_yolov11L.pt)
            conf: Confidence threshold (default 0.25)
            iou: IOU threshold for NMS (default 0.7, matches training)
        """
        self.conf = conf
        self.iou = iou
        self.model = None

        # Load YOLOv11 model directly from ultralytics
        try:
            if model_path is None:
                model_path = "models/weight5_multi_yolov11L.pt"

            self.model = YOLO(model_path)
            logger.info(f"✅ Detection pipeline initialized with YOLOv11 model: {model_path}")
            logger.info(f"   Inference size: {self.IMG_SIZE}x{self.IMG_SIZE}")
            logger.info(f"   Confidence threshold: {self.conf}")
            logger.info(f"   IOU threshold: {self.iou}")
        except Exception as e:
            logger.error(f"❌ Failed to load model: {e}")
            raise
    
    def preprocess(self, image: np.ndarray) -> np.ndarray:
        """
        Preprocess image with letterboxing (padding to maintain aspect ratio).
        Instead of stretching, adds black borders to fit 640x640.
        
        Args:
            image: Input image (numpy array in BGR format)
            
        Returns:
            Padded image (640x640) with black borders
        """
        h, w = image.shape[:2]
        
        # Calculate scale to fit within 640x640 while maintaining aspect ratio
        scale = min(self.IMG_SIZE / h, self.IMG_SIZE / w)
        new_h = int(h * scale)
        new_w = int(w * scale)
        
        # Resize image maintaining aspect ratio
        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        
        # Create 640x640 black canvas
        canvas = np.zeros((self.IMG_SIZE, self.IMG_SIZE, 3), dtype=np.uint8)
        
        # Calculate padding to center the image
        top = (self.IMG_SIZE - new_h) // 2
        left = (self.IMG_SIZE - new_w) // 2
        
        # Place resized image in center of canvas
        canvas[top:top+new_h, left:left+new_w] = resized
        
        return canvas
    
    def detect_from_file(self, image_path: str) -> Dict:
        """
        Detect ingredients from image file.
        
        Args:
            image_path: Path to image file
            
        Returns:
            Dictionary with detections
        """
        # Read image (auto-orient via cv2.imread)
        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"Image not found: {image_path}")
        
        # Store original size before preprocessing
        original_h, original_w = image.shape[:2]
        
        # Preprocess
        processed = self.preprocess(image)
        
        return self._run_detection(processed, original_size=(original_h, original_w))
    
    def detect_from_bytes(self, image_bytes: bytes) -> Dict:
        """
        Detect ingredients from image bytes.
        
        Args:
            image_bytes: Image as bytes
            
        Returns:
            Dictionary with detections
        """
        # Convert bytes to numpy array
        nparr = np.frombuffer(image_bytes, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if image is None:
            raise ValueError("Failed to decode image from bytes")
        
        # Store original size before preprocessing
        original_h, original_w = image.shape[:2]
        
        # Preprocess
        processed = self.preprocess(image)
        
        return self._run_detection(processed, original_size=(original_h, original_w))
    
    def detect_batch(self, image_paths: List[str]) -> List[Dict]:
        """
        Batch detection for multiple images (optimized for GPU).
        
        Args:
            image_paths: List of image file paths
            
        Returns:
            List of detection results
        """
        results = []
        
        try:
            for image_path in image_paths:
                result = self.detect_from_file(image_path)
                results.append(result)
                logger.info(f"Processed {image_path}: {result['count']} detections")
        except Exception as e:
            logger.error(f"❌ Batch detection failed: {e}")
            raise
        
        return results
    
    def _run_detection(self, image: np.ndarray, original_size: Tuple[int, int]) -> Dict:
        """
        Run detection inference on preprocessed image.
        
        Args:
            image: Preprocessed image (640x640)
            original_size: Original image size (h, w)
            
        Returns:
            Detection results dictionary
        """
        if self.model is None:
            raise RuntimeError("Model not loaded. Check model_path configuration.")
        
        try:
            logger.info(f"Running YOLOv11 inference on {self.IMG_SIZE}x{self.IMG_SIZE} image")
            
            # Run inference
            results = self.model(
                image,
                conf=self.conf,
                iou=self.iou,
                imgsz=self.IMG_SIZE,
                verbose=False
            )
            
            # Parse results (YOLOv11 format)
            detections = self._parse_results(results[0])
            logger.info(f"Detections found: {len(detections)}")
            
            return {
                'detections': detections,
                'raw_detections': detections,
                'image_size': original_size,
                'inference_size': (self.IMG_SIZE, self.IMG_SIZE),
                'count': len(detections)
            }
            
        except Exception as e:
            logger.error(f"❌ Detection failed: {e}")
            raise
    
    def _parse_results(self, results) -> List[Dict]:
        """
        Parse YOLOv11l Results object.
        
        Args:
            results: YOLOv11 Results object (from model.predict)
            
        Returns:
            List of detection dictionaries
        """
        detections = []
        
        try:
            # YOLOv11 format: results.boxes contains detection info
            if hasattr(results, 'boxes') and results.boxes is not None:
                boxes = results.boxes
                
                for i, box in enumerate(boxes):
                    # Extract coordinates (xyxy format: x1, y1, x2, y2)
                    xyxy = box.xyxy[0].cpu().numpy()
                    x1, y1, x2, y2 = xyxy.tolist()
                    
                    # Extract confidence
                    conf = float(box.conf[0])
                    
                    # Extract class
                    cls_id = int(box.cls[0])
                    cls_name = results.names[cls_id] if hasattr(results, 'names') else \
                               self.DEFAULT_CLASSES[cls_id % len(self.DEFAULT_CLASSES)]
                    
                    detections.append({
                        'bbox': [float(x1), float(y1), float(x2), float(y2)],
                        'confidence': conf,
                        'class_id': cls_id,
                        'class_name': cls_name,
                        'area': float((x2 - x1) * (y2 - y1))
                    })
                
                logger.info(f"Parsed {len(detections)} detections from YOLOv11 results")
            
            else:
                logger.warning("No boxes found in YOLOv11 results")
            
        except Exception as e:
            logger.warning(f"Result parsing issue: {e}")
        
        return detections
    
    def get_ingredient_list(self, detections: List[Dict]) -> List[str]:
        """
        Extract unique ingredient names.
        
        Args:
            detections: List of detection dictionaries
            
        Returns:
            List of unique ingredient names
        """
        ingredients = []
        seen = set()
        
        for detection in detections:
            name = detection['class_name']
            if name.lower() not in seen:
                ingredients.append(name)
                seen.add(name.lower())
        
        return ingredients