detection_engine.py

"""
Detection Engine
Handles YOLOv8 object detection and target matching
"""

import logging
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from ultralytics import YOLO
import time

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DetectionEngine:
    """YOLOv8-based object detection engine with target matching"""

    def __init__(self, model_path: str = "yolov8s.pt", confidence_threshold: float = 0.5):
        """
        Initialize detection engine

        Args:
            model_path: Path to YOLO model file
            confidence_threshold: Minimum confidence for detections
        """
        self.model_path = model_path
        self.confidence_threshold = confidence_threshold
        self.model: Optional[YOLO] = None
        self.class_names: List[str] = []
        self.total_detections = 0
        self.target_detections = 0

    def load_model(self) -> bool:
        """
        Load YOLO model

        Returns:
            True if model loaded successfully
        """
        try:
            logger.info(f"Loading YOLO model: {self.model_path}")
            self.model = YOLO(self.model_path)
            self.class_names = list(self.model.names.values())
            logger.info(f"Model loaded successfully. Classes: {len(self.class_names)}")
            return True
        except Exception as e:
            logger.error(f"Error loading YOLO model: {e}")
            # Try to load default model as fallback
            if self.model_path != "yolov8s.pt":
                logger.info("Attempting to load default model (yolov8s.pt) as fallback")
                try:
                    self.model = YOLO("yolov8s.pt")
                    self.class_names = list(self.model.names.values())
                    self.model_path = "yolov8s.pt"
                    logger.info("Default model loaded successfully")
                    return True
                except Exception as fallback_error:
                    logger.error(f"Failed to load fallback model: {fallback_error}")
            return False

    def detect(self, frame: np.ndarray, target_classes: Optional[List[str]] = None) -> List[Dict[str, Any]]:
        """
        Run detection on frame

        Args:
            frame: Input frame (BGR format)
            target_classes: Optional list of target class names to filter

        Returns:
            List of detection dictionaries
        """
        if self.model is None:
            logger.error("Model not loaded. Call load_model() first.")
            return []

        try:
            # Run inference
            results = self.model(frame, conf=self.confidence_threshold, verbose=False)

            detections = []

            # Process results
            for result in results:
                boxes = result.boxes

                for i, box in enumerate(boxes):
                    # Extract box data
                    xyxy = box.xyxy[0].cpu().numpy()  # Bounding box coordinates
                    conf = float(box.conf[0].cpu().numpy())  # Confidence
                    cls = int(box.cls[0].cpu().numpy())  # Class ID
                    class_name = self.model.names[cls]

                    # Check if this is a target class
                    if target_classes and class_name not in target_classes:
                        continue

                    detection = {
                        'class': class_name,
                        'class_id': cls,
                        'confidence': round(conf, 4),
                        'bbox': [int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])],
                        'frame_id': self.total_detections
                    }

                    detections.append(detection)
                    self.total_detections += 1

                    if target_classes and class_name in target_classes:
                        self.target_detections += 1

            return detections

        except Exception as e:
            logger.error(f"Error during detection: {e}")
            return []

    def get_class_names(self) -> List[str]:
        """Get list of all available class names"""
        return self.class_names

    def validate_target_class(self, target_class: str) -> bool:
        """
        Check if target class exists in model

        Args:
            target_class: Class name to validate

        Returns:
            True if class exists in model
        """
        return target_class.lower() in [name.lower() for name in self.class_names]

    def get_stats(self) -> Dict[str, int]:
        """Get detection statistics"""
        return {
            'total_detections': self.total_detections,
            'target_detections': self.target_detections
        }

    def reset_stats(self):
        """Reset detection counters"""
        self.total_detections = 0
        self.target_detections = 0