midlajvalappil/Real-time_Object_Detection_with_YOLO

1"""2YOLO Object Detection Module3Handles loading YOLO models and performing object detection on images.4"""5 6import cv27import numpy as np8import torch9from ultralytics import YOLO10from typing import List, Tuple, Dict, Any11import logging12 13# Configure logging14logging.basicConfig(level=logging.INFO)15logger = logging.getLogger(__name__)16 17class YOLODetector:18    """19    YOLO Object Detector class for real-time object detection.20    """21    22    def __init__(self, model_name: str = "yolov8n.pt", confidence_threshold: float = 0.5):23        """24        Initialize the YOLO detector.25        26        Args:27            model_name (str): Name of the YOLO model to use28            confidence_threshold (float): Minimum confidence threshold for detections29        """30        self.model_name = model_name31        self.confidence_threshold = confidence_threshold32        self.model = None33        self.device = "cuda" if torch.cuda.is_available() else "cpu"34        35        # COCO class names (80 classes)36        self.class_names = [37            'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',38            'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',39            'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',40            'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',41            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',42            'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',43            'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',44            'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',45            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',46            'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',47            'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',48            'toothbrush'49        ]50        51        self.load_model()52    53    def load_model(self) -> bool:54        """55        Load the YOLO model.56        57        Returns:58            bool: True if model loaded successfully, False otherwise59        """60        try:61            logger.info(f"Loading YOLO model: {self.model_name}")62            self.model = YOLO(self.model_name)63            self.model.to(self.device)64            logger.info(f"Model loaded successfully on device: {self.device}")65            return True66        except Exception as e:67            logger.error(f"Error loading model: {str(e)}")68            return False69    70    def detect_objects(self, image: np.ndarray) -> List[Dict[str, Any]]:71        """72        Perform object detection on an image.73        74        Args:75            image (np.ndarray): Input image in BGR format76            77        Returns:78            List[Dict[str, Any]]: List of detected objects with their properties79        """80        if self.model is None:81            logger.error("Model not loaded")82            return []83        84        try:85            # Run inference86            results = self.model(image, conf=self.confidence_threshold, verbose=False)87            88            detections = []89            for result in results:90                boxes = result.boxes91                if boxes is not None:92                    for box in boxes:93                        # Extract box coordinates94                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()95                        confidence = box.conf[0].cpu().numpy()96                        class_id = int(box.cls[0].cpu().numpy())97                        98                        # Get class name99                        class_name = self.class_names[class_id] if class_id < len(self.class_names) else f"class_{class_id}"100                        101                        detection = {102                            'bbox': [int(x1), int(y1), int(x2), int(y2)],103                            'confidence': float(confidence),104                            'class_id': class_id,105                            'class_name': class_name106                        }107                        detections.append(detection)108            109            return detections110            111        except Exception as e:112            logger.error(f"Error during detection: {str(e)}")113            return []114    115    def draw_detections(self, image: np.ndarray, detections: List[Dict[str, Any]]) -> np.ndarray:116        """117        Draw bounding boxes and labels on the image.118        119        Args:120            image (np.ndarray): Input image121            detections (List[Dict[str, Any]]): List of detections122            123        Returns:124            np.ndarray: Image with drawn detections125        """126        annotated_image = image.copy()127        128        for detection in detections:129            bbox = detection['bbox']130            confidence = detection['confidence']131            class_name = detection['class_name']132            133            x1, y1, x2, y2 = bbox134            135            # Draw bounding box136            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)137            138            # Prepare label139            label = f"{class_name}: {confidence:.2f}"140            141            # Get text size for background rectangle142            (text_width, text_height), baseline = cv2.getTextSize(143                label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1144            )145            146            # Draw background rectangle for text147            cv2.rectangle(148                annotated_image,149                (x1, y1 - text_height - baseline - 5),150                (x1 + text_width, y1),151                (0, 255, 0),152                -1153            )154            155            # Draw text156            cv2.putText(157                annotated_image,158                label,159                (x1, y1 - baseline - 2),160                cv2.FONT_HERSHEY_SIMPLEX,161                0.5,162                (0, 0, 0),163                1164            )165        166        return annotated_image167    168    def get_detection_stats(self, detections: List[Dict[str, Any]]) -> Dict[str, Any]:169        """170        Get statistics about the detections.171        172        Args:173            detections (List[Dict[str, Any]]): List of detections174            175        Returns:176            Dict[str, Any]: Detection statistics177        """178        if not detections:179            return {180                'total_objects': 0,181                'class_counts': {},182                'avg_confidence': 0.0,183                'max_confidence': 0.0,184                'min_confidence': 0.0185            }186        187        class_counts = {}188        confidences = []189        190        for detection in detections:191            class_name = detection['class_name']192            confidence = detection['confidence']193            194            class_counts[class_name] = class_counts.get(class_name, 0) + 1195            confidences.append(confidence)196        197        return {198            'total_objects': len(detections),199            'class_counts': class_counts,200            'avg_confidence': np.mean(confidences),201            'max_confidence': np.max(confidences),202            'min_confidence': np.min(confidences)203        }204    205    def update_confidence_threshold(self, threshold: float):206        """207        Update the confidence threshold.208        209        Args:210            threshold (float): New confidence threshold211        """212        self.confidence_threshold = max(0.0, min(1.0, threshold))213        logger.info(f"Confidence threshold updated to: {self.confidence_threshold}")214