detector/detector.py

"""
Object detection using YOLOv8n (nano). CPU-friendly, no face recognition.
Returns bounding boxes + labels for person, dog, cat, etc.
"""

from dataclasses import dataclass
from typing import List, Tuple

import cv2
from ultralytics import YOLO


# Classes we care about for door / animal use case
PERSON_LABEL = "person"
ANIMAL_LABELS = frozenset({"dog", "cat", "bird", "horse", "sheep", "cow", "bear"})


@dataclass
class Detection:
    """Single detection: label + bbox (xyxy)."""
    label: str
    confidence: float
    xyxy: Tuple[float, float, float, float]  # x1, y1, x2, y2


class Detector:
    """
    YOLOv8n-based detector. Load once, run on frames.
    """

    def __init__(self, model_path: str = "yolov8n.pt"):
        self.model = YOLO(model_path)
        self.names = self.model.names

    def detect(self, frame, conf_threshold: float = 0.25) -> List[Detection]:
        """
        Run detection on a BGR frame (e.g. from cv2).
        Returns list of Detection with label, confidence, bbox.
        """
        results = self.model(frame, conf=conf_threshold, verbose=False)
        out = []
        for r in results:
            if r.boxes is None:
                continue
            for box in r.boxes:
                cls_id = int(box.cls[0])
                label = self.names[cls_id]
                conf = float(box.conf[0])
                xyxy = tuple(map(float, box.xyxy[0]))
                out.append(Detection(label=label, confidence=conf, xyxy=xyxy))
        return out

    def is_person(self, label: str) -> bool:
        return label == PERSON_LABEL

    def is_animal(self, label: str) -> bool:
        return label in ANIMAL_LABELS

    def annotate_frame(self, frame, detections: List[Detection]):
        """Draw bounding boxes and labels on a copy of the frame."""
        annotated = frame.copy()
        for d in detections:
            x1, y1, x2, y2 = map(int, d.xyxy)
            cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 2)
            text = "{} {:.2f}".format(d.label, d.confidence)
            cv2.putText(
                annotated, text, (x1, y1 - 8),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA
            )
        return annotated