File size: 4,445 Bytes
e1c010e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """
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
|