Spaces:
Sleeping
Sleeping
File size: 8,314 Bytes
3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c 3c82458 3415c7c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 |
"""
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 |