detection.py

"""
Enhanced Detection with Better Confidence Filtering and NMS (Enhancement 7)
"""
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:
    """Detection data structure"""
    bbox: List[float]  # [x1, y1, x2, y2]
    confidence: float
    image_crop: Optional[np.ndarray] = None

class EnhancedDogDetector:
    """
    Enhanced YOLOv8 detector with improved filtering (Enhancement 7)
    """
    def __init__(self,
                 confidence_threshold: float = 0.50,  # Increased from 0.45
                 nms_threshold: float = 0.4,  # Non-maximum suppression
                 device: str = 'cuda'):
        """
        Initialize detector with enhanced filtering
        Args:
            confidence_threshold: Higher threshold reduces false positives
            nms_threshold: Lower = stricter NMS, removes more overlapping boxes
            device: 'cuda' or 'cpu'
        """
        self.confidence_threshold = confidence_threshold
        self.nms_threshold = nms_threshold
        self.device = device if torch.cuda.is_available() else 'cpu'
        
        # Load YOLOv8 medium model
        self.model = YOLO('yolov8m.pt')
        self.model.to(self.device)
        
        # COCO class ID for dog
        self.dog_class_id = 16
        
        # ENHANCEMENT 7: Size constraints
        self.min_detection_area = 900  # 30x30 pixels minimum
        self.max_detection_area = 640000  # 800x800 pixels maximum
        
        print(f"✅ Enhanced Detector initialized")
        print(f"  Confidence: {self.confidence_threshold:.2f}")
        print(f"  NMS threshold: {self.nms_threshold:.2f}")
        print(f"  Min area: {self.min_detection_area}px²")
    
    def _apply_custom_nms(self, boxes, scores, iou_threshold=0.4):
        """
        ENHANCEMENT 7: Custom NMS for better duplicate removal
        """
        if len(boxes) == 0:
            return []
        
        # Convert to numpy arrays
        boxes = np.array(boxes)
        scores = np.array(scores)
        
        # Get coordinates
        x1 = boxes[:, 0]
        y1 = boxes[:, 1]
        x2 = boxes[:, 2]
        y2 = boxes[:, 3]
        
        # Compute areas
        areas = (x2 - x1) * (y2 - y1)
        
        # Sort by score
        order = scores.argsort()[::-1]
        
        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)
            
            # Compute IoU with remaining boxes
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])
            
            w = np.maximum(0.0, xx2 - xx1)
            h = np.maximum(0.0, yy2 - yy1)
            inter = w * h
            
            iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
            
            # Keep boxes with IoU less than threshold
            inds = np.where(iou <= iou_threshold)[0]
            order = order[inds + 1]
        
        return keep
    
    def _filter_by_size(self, detections: List[Detection]) -> List[Detection]:
        """
        ENHANCEMENT 7: Size-based filtering to remove false positives
        """
        filtered = []
        
        for det in detections:
            width = det.bbox[2] - det.bbox[0]
            height = det.bbox[3] - det.bbox[1]
            area = width * height
            
            # Check area constraints
            if area < self.min_detection_area or area > self.max_detection_area:
                continue
            
            # Check aspect ratio (dogs shouldn't be extreme shapes)
            if width > 0 and height > 0:
                aspect_ratio = width / height
                if aspect_ratio < 0.2 or aspect_ratio > 5.0:
                    continue
            
            filtered.append(det)
        
        return filtered
    
    def _filter_by_confidence_quality(self, detections: List[Detection]) -> List[Detection]:
        """
        ENHANCEMENT 7: Advanced confidence filtering
        """
        if not detections:
            return []
        
        # Calculate confidence statistics
        confidences = [d.confidence for d in detections]
        mean_conf = np.mean(confidences)
        std_conf = np.std(confidences)
        
        filtered = []
        for det in detections:
            # Base threshold
            if det.confidence < self.confidence_threshold:
                continue
            
            # Adaptive threshold: if confidence is much lower than mean, reject
            if len(detections) > 3:
                if det.confidence < mean_conf - std_conf:
                    continue
            
            filtered.append(det)
        
        return filtered
    
    def detect(self, frame: np.ndarray) -> List[Detection]:
        """
        Detect dogs with enhanced filtering
        """
        # Run YOLO inference
        results = self.model(frame,
                           conf=self.confidence_threshold * 0.9,  # Slightly lower for YOLO
                           classes=[self.dog_class_id],
                           verbose=False)
        
        initial_detections = []
        
        if results and len(results) > 0:
            result = results[0]
            
            if result.boxes is not None:
                boxes = result.boxes
                
                # Collect all boxes first
                all_boxes = []
                all_scores = []
                
                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)
                    
                    # Ensure valid coordinates
                    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
                    
                    all_boxes.append([x1, y1, x2, y2])
                    all_scores.append(float(boxes.conf[i]))
                
                # ENHANCEMENT 7: Apply custom NMS
                if all_boxes:
                    keep_indices = self._apply_custom_nms(
                        all_boxes, 
                        all_scores, 
                        iou_threshold=self.nms_threshold
                    )
                    
                    # Create detections for kept boxes
                    for idx in keep_indices:
                        bbox = all_boxes[idx]
                        conf = all_scores[idx]
                        
                        # Crop dog image
                        x1, y1, x2, y2 = bbox
                        dog_crop = frame[y1:y2, x1:x2].copy()
                        
                        detection = Detection(
                            bbox=bbox,
                            confidence=conf,
                            image_crop=dog_crop
                        )
                        initial_detections.append(detection)
        
        # ENHANCEMENT 7: Apply additional filters
        filtered_detections = self._filter_by_size(initial_detections)
        filtered_detections = self._filter_by_confidence_quality(filtered_detections)
        
        # Debug info
        if len(initial_detections) != len(filtered_detections):
            print(f"  🔍 Detection filter: {len(initial_detections)} → {len(filtered_detections)}")
        
        return filtered_detections
    
    def set_confidence(self, threshold: float):
        """Update detection confidence threshold"""
        self.confidence_threshold = max(0.1, min(1.0, threshold))
        print(f"Detection confidence updated: {self.confidence_threshold:.2f}")
    
    def set_nms_threshold(self, threshold: float):
        """Update NMS threshold"""
        self.nms_threshold = max(0.1, min(0.9, threshold))
        print(f"NMS threshold updated: {self.nms_threshold:.2f}")

# Compatibility alias
DogDetector = EnhancedDogDetector