Create detection.py

detection.py

@@ -0,0 +1,95 @@
+import cv2
+import numpy as np
+import torch
+from ultralytics import YOLO
+from typing import List, Tuple, Optional
+from dataclasses import dataclass
+
+@dataclass
+class Detection:
+    """Simple detection data structure"""
+    bbox: List[float]  # [x1, y1, x2, y2]
+    confidence: float
+    image_crop: Optional[np.ndarray] = None  # Cropped dog image
+    
+class DogDetector:
+    """
+    Simplified YOLOv8 detector optimized for dogs
+    Uses standard pretrained model - no custom training needed
+    """
+    
+    def __init__(self, 
+                 confidence_threshold: float = 0.45,
+                 device: str = 'cuda'):
+        """
+        Initialize detector
+        
+        Args:
+            confidence_threshold: Min confidence for detections (0.45 works well)
+            device: 'cuda' for GPU, 'cpu' for CPU
+        """
+        self.confidence_threshold = confidence_threshold
+        self.device = device if torch.cuda.is_available() else 'cpu'
+        
+        # Load YOLOv8 medium model (good balance of speed/accuracy)
+        self.model = YOLO('yolov8m.pt')
+        self.model.to(self.device)
+        
+        # COCO class ID for dog
+        self.dog_class_id = 16
+        
+    def detect(self, frame: np.ndarray) -> List[Detection]:
+        """
+        Detect dogs in frame
+        
+        Args:
+            frame: BGR image from OpenCV
+            
+        Returns:
+            List of Detection objects with bounding boxes and crops
+        """
+        # Run YOLO inference
+        results = self.model(frame, 
+                            conf=self.confidence_threshold,
+                            classes=[self.dog_class_id],  # Only detect dogs
+                            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)):
+                    # Get bbox coordinates
+                    x1, y1, x2, y2 = boxes.xyxy[i].cpu().numpy()
+                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+                    
+                    # Ensure valid coordinates
+                    h, w = frame.shape[:2]
+                    x1 = max(0, x1)
+                    y1 = max(0, y1)
+                    x2 = min(w, x2)
+                    y2 = min(h, y2)
+                    
+                    # Skip invalid boxes
+                    if x2 <= x1 or y2 <= y1:
+                        continue
+                    
+                    # Crop dog image
+                    dog_crop = frame[y1:y2, x1:x2].copy()
+                    
+                    # Create detection
+                    detection = Detection(
+                        bbox=[x1, y1, x2, y2],
+                        confidence=float(boxes.conf[i]),
+                        image_crop=dog_crop
+                    )
+                    detections.append(detection)
+        
+        return detections
+    
+    def set_confidence(self, threshold: float):
+        """Update detection confidence threshold"""
+        self.confidence_threshold = max(0.1, min(1.0, threshold))