| """ |
| Vision Module — Object detection using YOLOv8. |
| |
| Wraps the Ultralytics YOLO model to detect objects in images and return |
| structured results. The model is loaded **once** in the constructor and |
| reused for all subsequent detections. |
| """ |
|
|
| import logging |
| from typing import Any |
|
|
| import numpy as np |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class ObjectDetector: |
| """ |
| Detects objects in images using a pretrained YOLOv8 model. |
| |
| Attributes: |
| model: The loaded YOLO model instance (cached after first load). |
| """ |
|
|
| def __init__(self, model_name: str = "yolov8n.pt", confidence: float = 0.25): |
| """ |
| Initialize the detector and load the YOLO model. |
| |
| Args: |
| model_name: Name or path of the YOLO model weights. |
| confidence: Minimum confidence threshold for detections. |
| """ |
| self.model_name = model_name |
| self.confidence = confidence |
| self.model = None |
|
|
| self._load_model() |
|
|
| def _load_model(self) -> None: |
| """Load the YOLO model. Called once during initialization.""" |
| try: |
| from ultralytics import YOLO |
|
|
| logger.info("Loading YOLO model: %s", self.model_name) |
| self.model = YOLO(self.model_name) |
| logger.info("YOLO model loaded successfully.") |
| except Exception as exc: |
| logger.error("Failed to load YOLO model '%s': %s", self.model_name, exc) |
| raise RuntimeError( |
| f"Could not load YOLO model '{self.model_name}'. " |
| f"Ensure ultralytics is installed and the model file is accessible." |
| ) from exc |
|
|
| def detect(self, image: np.ndarray) -> list[dict[str, Any]]: |
| """ |
| Run object detection on an image. |
| |
| Args: |
| image: Input image as a NumPy array (H, W, C) in BGR or RGB format. |
| |
| Returns: |
| List of detected objects, each containing: |
| - ``label`` (str): Class name. |
| - ``confidence`` (float): Detection confidence (0-1). |
| - ``bbox`` (list[int]): Bounding box as [x1, y1, x2, y2]. |
| |
| Raises: |
| ValueError: If the input image is invalid. |
| RuntimeError: If the model is not loaded. |
| """ |
| if self.model is None: |
| raise RuntimeError("YOLO model is not loaded.") |
|
|
| if image is None or not isinstance(image, np.ndarray): |
| raise ValueError("Invalid image: expected a NumPy ndarray.") |
|
|
| if image.ndim < 2: |
| raise ValueError( |
| f"Invalid image dimensions: expected 2D or 3D array, " |
| f"got {image.ndim}D." |
| ) |
|
|
| logger.info( |
| "Running detection on image of shape %s with confidence=%.2f", |
| image.shape, |
| self.confidence, |
| ) |
|
|
| try: |
| results = self.model(image, conf=self.confidence, verbose=False) |
| except Exception as exc: |
| logger.error("YOLO inference failed: %s", exc) |
| return [] |
|
|
| detections = self._parse_results(results) |
| logger.info("Detected %d objects.", len(detections)) |
| return detections |
|
|
| def _parse_results(self, results: Any) -> list[dict[str, Any]]: |
| """ |
| Parse YOLO results into a structured list of detections. |
| |
| Args: |
| results: Raw YOLO results object. |
| |
| Returns: |
| Structured list of detection dictionaries. |
| """ |
| detections: list[dict[str, Any]] = [] |
|
|
| for result in results: |
| boxes = result.boxes |
| if boxes is None: |
| continue |
|
|
| for i in range(len(boxes)): |
| try: |
| bbox = boxes.xyxy[i].cpu().numpy().tolist() |
| confidence = float(boxes.conf[i].cpu().numpy()) |
| class_id = int(boxes.cls[i].cpu().numpy()) |
| label = result.names.get(class_id, f"class_{class_id}") |
|
|
| detections.append( |
| { |
| "label": label, |
| "confidence": round(confidence, 4), |
| "bbox": [int(coord) for coord in bbox], |
| } |
| ) |
| except (IndexError, KeyError, AttributeError) as exc: |
| logger.warning("Failed to parse detection %d: %s", i, exc) |
| continue |
|
|
| return detections |
|
|
