Spaces:
Sleeping
Sleeping
File size: 4,193 Bytes
8ffc27e | 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 | """
Object Detection Engine using YOLO
"""
import cv2
import numpy as np
from ultralytics import YOLO
import os
class ObjectDetector:
def __init__(self, model_path='runs/detect/safety_detector/weights/best.pt'):
"""
Initialize the YOLO detector
"""
print(f"🔍 Loading model: {model_path}")
# Check if model exists, use pretrained if not
if not os.path.exists(model_path):
print("⚠️ Trained model not found. Using pretrained YOLO11...")
model_path = 'yolo11n.pt'
self.model = YOLO(model_path)
# Class names (from Hard Hat Workers dataset)
self.class_names = {
0: 'Helmet',
1: 'No-Helmet',
2: 'Vest',
3: 'Person'
}
# Colors for each class (BGR)
self.colors = {
'Helmet': (0, 255, 0), # Green
'No-Helmet': (0, 0, 255), # Red
'Vest': (255, 200, 0), # Yellow
'Person': (255, 0, 0) # Blue
}
print(f"✅ Detector initialized with {len(self.class_names)} classes")
def detect_frame(self, frame):
"""
Detect objects in a single frame
"""
# Run inference
results = self.model(frame, verbose=False)
# Extract detections
detections = []
for result in results:
boxes = result.boxes
if boxes is not None:
for box in boxes:
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
cls = int(box.cls[0])
conf = float(box.conf[0])
# Only include confident detections
if conf > 0.3:
detections.append({
'bbox': [x1, y1, x2, y2],
'class': self.class_names.get(cls, 'Unknown'),
'class_id': cls,
'confidence': conf
})
return detections
def draw_detections(self, frame, detections):
"""
Draw bounding boxes and labels on frame
"""
annotated_frame = frame.copy()
for det in detections:
x1, y1, x2, y2 = det['bbox']
cls_name = det['class']
conf = det['confidence']
color = self.colors.get(cls_name, (255, 255, 255))
# Draw rectangle
cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
# Draw label with background
label = f"{cls_name} {conf:.2f}"
(label_w, label_h), _ = cv2.getTextSize(
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
)
cv2.rectangle(
annotated_frame,
(x1, y1 - label_h - 8),
(x1 + label_w + 8, y1),
color,
-1
)
cv2.putText(
annotated_frame,
label,
(x1 + 4, y1 - 4),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(255, 255, 255),
2
)
return annotated_frame
def process_image(self, image_path):
"""
Process an image file
"""
frame = cv2.imread(image_path)
if frame is None:
raise ValueError(f"Could not read image: {image_path}")
detections = self.detect_frame(frame)
annotated = self.draw_detections(frame, detections)
return annotated, detections
def get_stats(self, detections):
"""
Get statistics from detections
"""
stats = {}
for det in detections:
cls_name = det['class']
stats[cls_name] = stats.get(cls_name, 0) + 1
return stats
# Test the detector
if __name__ == "__main__":
print("🧪 Testing ObjectDetector...")
detector = ObjectDetector()
print("✅ ObjectDetector initialized successfully!")
|