AryanS17/Computer-Vision-Accessibility-Tool

1"""2detector.py3-----------4Thin wrapper around Ultralytics YOLO for real-time object detection.5Kept separate from the Streamlit app so it's easy to swap models6(YOLOv8n for speed on CPU, YOLOv8m/YOLOv10 for accuracy on GPU) or7unit-test independently.8"""9 10from typing import List, Dict11import numpy as np12from ultralytics import YOLO13 14 15class ObjectDetector:16    def __init__(self, model_path: str = "yolov8n.pt", confidence_threshold: float = 0.45):17        """18        model_path: 'yolov8n.pt' (nano, fastest, good default for CPU/webcam).19                    Swap to 'yolov8s.pt' or 'yolov10n.pt' for a different20                    speed/accuracy tradeoff. Ultralytics auto-downloads21                    weights on first run.22        """23        self.model = YOLO(model_path)24        self.confidence_threshold = confidence_threshold25        self.class_names = self.model.names  # id -> label string26 27    def detect(self, frame: np.ndarray) -> List[Dict]:28        """29        Runs inference on a single BGR frame (as returned by cv2.VideoCapture).30        Returns a list of dicts: {"label": str, "confidence": float, "box": (x1,y1,x2,y2)}31        """32        results = self.model.predict(33            source=frame,34            conf=self.confidence_threshold,35            verbose=False,36        )37 38        detections = []39        if not results:40            return detections41 42        result = results[0]43        boxes = result.boxes44        if boxes is None:45            return detections46 47        for box in boxes:48            cls_id = int(box.cls[0])49            label = self.class_names.get(cls_id, str(cls_id))50            confidence = float(box.conf[0])51            x1, y1, x2, y2 = box.xyxy[0].tolist()52            detections.append(53                {54                    "label": label,55                    "confidence": confidence,56                    "box": (x1, y1, x2, y2),57                }58            )59 60        return detections61 62    def annotate(self, frame: np.ndarray) -> np.ndarray:63        """Returns a copy of the frame with bounding boxes drawn (for the live preview)."""64        results = self.model.predict(source=frame, conf=self.confidence_threshold, verbose=False)65        if not results:66            return frame67        return results[0].plot()  # ultralytics built-in annotator (BGR np.ndarray)68