|
|
""" |
|
|
detection.py - Dog Detection using YOLOv8 |
|
|
""" |
|
|
|
|
|
import cv2 |
|
|
import numpy as np |
|
|
import torch |
|
|
from ultralytics import YOLO |
|
|
from typing import List, Optional |
|
|
from dataclasses import dataclass |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class Detection: |
|
|
bbox: List[float] |
|
|
confidence: float |
|
|
image_crop: Optional[np.ndarray] = None |
|
|
|
|
|
|
|
|
class DogDetector: |
|
|
def __init__(self, confidence_threshold: float = 0.50, device: str = 'cuda'): |
|
|
self.confidence_threshold = confidence_threshold |
|
|
self.device = device if torch.cuda.is_available() else 'cpu' |
|
|
self.model = YOLO('yolov8m.pt') |
|
|
self.model.to(self.device) |
|
|
self.dog_class_id = 16 |
|
|
self.min_detection_area = 900 |
|
|
|
|
|
print(f"✅ Dog Detector initialized on {self.device}") |
|
|
print(f" Confidence threshold: {self.confidence_threshold:.2f}") |
|
|
|
|
|
def detect(self, frame: np.ndarray) -> List[Detection]: |
|
|
if frame is None or frame.size == 0: |
|
|
return [] |
|
|
|
|
|
results = self.model( |
|
|
frame, |
|
|
conf=self.confidence_threshold, |
|
|
classes=[self.dog_class_id], |
|
|
verbose=False |
|
|
) |
|
|
|
|
|
detections = [] |
|
|
|
|
|
if results and len(results) > 0: |
|
|
result = results[0] |
|
|
|
|
|
if result.boxes is not None: |
|
|
boxes = result.boxes |
|
|
|
|
|
for i in range(len(boxes)): |
|
|
x1, y1, x2, y2 = boxes.xyxy[i].cpu().numpy() |
|
|
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) |
|
|
|
|
|
h, w = frame.shape[:2] |
|
|
x1 = max(0, min(w-1, x1)) |
|
|
y1 = max(0, min(h-1, y1)) |
|
|
x2 = max(0, min(w, x2)) |
|
|
y2 = max(0, min(h, y2)) |
|
|
|
|
|
if x2 <= x1 or y2 <= y1: |
|
|
continue |
|
|
|
|
|
area = (x2 - x1) * (y2 - y1) |
|
|
if area < self.min_detection_area: |
|
|
continue |
|
|
|
|
|
confidence = float(boxes.conf[i]) |
|
|
dog_crop = frame[y1:y2, x1:x2].copy() |
|
|
|
|
|
detection = Detection( |
|
|
bbox=[x1, y1, x2, y2], |
|
|
confidence=confidence, |
|
|
image_crop=dog_crop |
|
|
) |
|
|
|
|
|
detections.append(detection) |
|
|
|
|
|
return detections |
|
|
|
|
|
def set_confidence(self, threshold: float): |
|
|
self.confidence_threshold = max(0.1, min(1.0, threshold)) |
|
|
print(f"Detection confidence updated: {self.confidence_threshold:.2f}") |
|
|
|
|
|
|
|
|
EnhancedDogDetector = DogDetector |
