SahilCodevally/codevally-vision-language-action

1"""2Vision Module — Object detection using YOLOv8.3 4Wraps the Ultralytics YOLO model to detect objects in images and return5structured results. The model is loaded **once** in the constructor and6reused for all subsequent detections.7"""8 9import logging10from typing import Any11 12import numpy as np13 14logger = logging.getLogger(__name__)15 16 17class ObjectDetector:18    """19    Detects objects in images using a pretrained YOLOv8 model.20 21    Attributes:22        model: The loaded YOLO model instance (cached after first load).23    """24 25    def __init__(self, model_name: str = "yolov8n.pt", confidence: float = 0.25):26        """27        Initialize the detector and load the YOLO model.28 29        Args:30            model_name: Name or path of the YOLO model weights.31            confidence: Minimum confidence threshold for detections.32        """33        self.model_name = model_name34        self.confidence = confidence35        self.model = None36 37        self._load_model()38 39    def _load_model(self) -> None:40        """Load the YOLO model. Called once during initialization."""41        try:42            from ultralytics import YOLO43 44            logger.info("Loading YOLO model: %s", self.model_name)45            self.model = YOLO(self.model_name)46            logger.info("YOLO model loaded successfully.")47        except Exception as exc:48            logger.error("Failed to load YOLO model '%s': %s", self.model_name, exc)49            raise RuntimeError(50                f"Could not load YOLO model '{self.model_name}'. "51                f"Ensure ultralytics is installed and the model file is accessible."52            ) from exc53 54    def detect(self, image: np.ndarray) -> list[dict[str, Any]]:55        """56        Run object detection on an image.57 58        Args:59            image: Input image as a NumPy array (H, W, C) in BGR or RGB format.60 61        Returns:62            List of detected objects, each containing:63                - ``label`` (str): Class name.64                - ``confidence`` (float): Detection confidence (0-1).65                - ``bbox`` (list[int]): Bounding box as [x1, y1, x2, y2].66 67        Raises:68            ValueError: If the input image is invalid.69            RuntimeError: If the model is not loaded.70        """71        if self.model is None:72            raise RuntimeError("YOLO model is not loaded.")73 74        if image is None or not isinstance(image, np.ndarray):75            raise ValueError("Invalid image: expected a NumPy ndarray.")76 77        if image.ndim < 2:78            raise ValueError(79                f"Invalid image dimensions: expected 2D or 3D array, "80                f"got {image.ndim}D."81            )82 83        logger.info(84            "Running detection on image of shape %s with confidence=%.2f",85            image.shape,86            self.confidence,87        )88 89        try:90            results = self.model(image, conf=self.confidence, verbose=False)91        except Exception as exc:92            logger.error("YOLO inference failed: %s", exc)93            return []94 95        detections = self._parse_results(results)96        logger.info("Detected %d objects.", len(detections))97        return detections98 99    def _parse_results(self, results: Any) -> list[dict[str, Any]]:100        """101        Parse YOLO results into a structured list of detections.102 103        Args:104            results: Raw YOLO results object.105 106        Returns:107            Structured list of detection dictionaries.108        """109        detections: list[dict[str, Any]] = []110 111        for result in results:112            boxes = result.boxes113            if boxes is None:114                continue115 116            for i in range(len(boxes)):117                try:118                    bbox = boxes.xyxy[i].cpu().numpy().tolist()119                    confidence = float(boxes.conf[i].cpu().numpy())120                    class_id = int(boxes.cls[i].cpu().numpy())121                    label = result.names.get(class_id, f"class_{class_id}")122 123                    detections.append(124                        {125                            "label": label,126                            "confidence": round(confidence, 4),127                            "bbox": [int(coord) for coord in bbox],128                        }129                    )130                except (IndexError, KeyError, AttributeError) as exc:131                    logger.warning("Failed to parse detection %d: %s", i, exc)132                    continue133 134        return detections135