""" 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