Spaces:
Runtime error
Runtime error
File size: 3,569 Bytes
0469d65 |
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 |
#!/usr/bin/env python3
"""
Test script for improved SafetyMaster Pro detection with stricter confidence thresholds
"""
import cv2
import numpy as np
from safety_detector import SafetyDetector
import time
def test_improved_detection():
"""Test the improved detection logic with stricter thresholds."""
print("🧪 Testing SafetyMaster Pro - Improved Detection Logic")
print("=" * 60)
# Initialize detector
detector = SafetyDetector()
print(f"\n📊 Confidence Thresholds:")
for equipment, threshold in detector.equipment_confidence_thresholds.items():
print(f" {equipment}: {threshold}")
print(f"\n🎯 Available Model Classes:")
classes = detector.get_model_classes()
for i, cls in enumerate(classes):
print(f" {i}: {cls}")
# Test with webcam
print(f"\n📹 Starting webcam test...")
print(" Press 'q' to quit")
print(" Press 's' to save current frame")
print(" Watch for improved accuracy with stricter thresholds")
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("❌ Error: Could not open webcam")
return
frame_count = 0
total_processing_time = 0
while True:
ret, frame = cap.read()
if not ret:
print("❌ Error: Could not read frame")
break
frame_count += 1
# Run detection
start_time = time.time()
results = detector.detect_safety_violations(frame)
processing_time = time.time() - start_time
total_processing_time += processing_time
# Draw detections
annotated_frame = detector.draw_detections(frame, results)
# Add performance info
avg_fps = frame_count / total_processing_time if total_processing_time > 0 else 0
cv2.putText(annotated_frame, f"Avg FPS: {avg_fps:.1f}",
(10, annotated_frame.shape[0] - 60),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
cv2.putText(annotated_frame, f"Frame: {frame_count}",
(10, annotated_frame.shape[0] - 30),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
# Print detection summary every 30 frames
if frame_count % 30 == 0:
print(f"\n📊 Frame {frame_count} Summary:")
print(f" People: {results['people_count']}")
print(f" Equipment detected: {results['safety_equipment']}")
print(f" Violations: {len(results['violations'])}")
if results['violations']:
for violation in results['violations']:
print(f" - {violation['description']}")
print(f" Processing time: {processing_time:.3f}s")
# Display frame
cv2.imshow('SafetyMaster Pro - Improved Detection', annotated_frame)
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
break
elif key == ord('s'):
# Save current frame
timestamp = time.strftime("%Y%m%d_%H%M%S")
filename = f"test_detection_{timestamp}.jpg"
cv2.imwrite(filename, annotated_frame)
print(f"💾 Saved frame as {filename}")
cap.release()
cv2.destroyAllWindows()
print(f"\n✅ Test completed!")
print(f" Total frames processed: {frame_count}")
print(f" Average FPS: {frame_count / total_processing_time:.1f}")
if __name__ == "__main__":
test_improved_detection() |