Update detection.py

detection.py

@@ -1,3 +1,6 @@
+"""
+Enhanced Detection with Better Confidence Filtering and NMS (Enhancement 7)
+"""
 import cv2
 import numpy as np
 import torch
@@ -7,89 +10,226 @@ from dataclasses import dataclass
 
 @dataclass
 class Detection:
-    """Simple detection data structure"""
+    """Detection data structure"""
     bbox: List[float]  # [x1, y1, x2, y2]
     confidence: float
-    image_crop: Optional[np.ndarray] = None  # Cropped dog image
-    
-class DogDetector:
+    image_crop: Optional[np.ndarray] = None
+
+class EnhancedDogDetector:
     """
-    Simplified YOLOv8 detector optimized for dogs
-    Uses standard pretrained model - no custom training needed
+    Enhanced YOLOv8 detector with improved filtering (Enhancement 7)
     """
-    
-    def __init__(self, 
-                 confidence_threshold: float = 0.45,
+    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
-        
+        Initialize detector with enhanced filtering
         Args:
-            confidence_threshold: Min confidence for detections (0.45 works well)
-            device: 'cuda' for GPU, 'cpu' for CPU
+            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 (good balance of speed/accuracy)
+        # Load YOLOv8 medium model
         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]:
+        # 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
         """
-        Detect dogs in frame
+        if len(boxes) == 0:
+            return []
         
-        Args:
-            frame: BGR image from OpenCV
+        # 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
             
-        Returns:
-            List of Detection objects with bounding boxes and crops
+            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,
-                            classes=[self.dog_class_id],  # Only detect dogs
-                            verbose=False)
+        results = self.model(frame,
+                           conf=self.confidence_threshold * 0.9,  # Slightly lower for YOLO
+                           classes=[self.dog_class_id],
+                           verbose=False)
         
-        detections = []
+        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)):
-                    # 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)
+                    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))
                     
-                    # 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
+                    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
                     )
-                    detections.append(detection)
+                    
+                    # 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 detections
+        return filtered_detections
     
     def set_confidence(self, threshold: float):
         """Update detection confidence threshold"""
-        self.confidence_threshold = max(0.1, min(1.0, 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