Spaces:
Runtime error
Runtime error
File size: 3,552 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 |
#!/usr/bin/env python3
"""
Test script for mask detection and violation logic
"""
from safety_detector import SafetyDetector
import cv2
import numpy as np
def test_violation_logic():
"""Test the violation detection logic."""
print("🧪 Testing SafetyMaster Pro Mask Detection & Violation Logic")
print("=" * 60)
# Initialize detector
detector = SafetyDetector()
# Test 1: Empty frame (no people, no equipment)
print("\n📋 Test 1: Empty frame")
empty_frame = np.zeros((480, 640, 3), dtype=np.uint8)
results = detector.detect_safety_violations(empty_frame)
print(f" People: {results['people_count']}")
print(f" Violations: {len(results['violations'])}")
print(f" Expected: 0 people, 0 violations ✅")
# Test 2: Check model classes
print("\n📋 Test 2: Model Classes")
classes = detector.get_model_classes()
mask_classes = [cls for cls in classes if 'mask' in cls.lower()]
print(f" Total classes: {len(classes)}")
print(f" Mask-related classes: {mask_classes}")
print(f" Expected: ['Mask', 'NO-Mask'] ✅")
# Test 3: Check PPE class mappings
print("\n📋 Test 3: PPE Class Mappings")
for category, variations in detector.ppe_classes.items():
if 'mask' in category:
print(f" {category}: {variations}")
# Test 4: Test with webcam (if available)
print("\n📋 Test 4: Live Camera Test")
print(" Starting webcam test - press 'q' to quit")
print(" Look for:")
print(" - Blue boxes around masks")
print(" - Red boxes around people without PPE")
print(" - Violation indicators on non-compliant people")
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print(" ❌ Could not open webcam")
return
frame_count = 0
while True:
ret, frame = cap.read()
if not ret:
break
# Process frame
results = detector.detect_safety_violations(frame)
annotated_frame = detector.draw_detections(frame, results)
# Add test info overlay
cv2.putText(annotated_frame, "SafetyMaster Pro - Mask Detection Test",
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
cv2.putText(annotated_frame, "Press 'q' to quit",
(10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
# Show detection info every 30 frames
if frame_count % 30 == 0:
print(f"\n Frame {frame_count}:")
print(f" - People detected: {results['people_count']}")
print(f" - Violations: {len(results['violations'])}")
print(f" - Equipment: {results['safety_equipment']}")
if results['violations']:
for violation in results['violations']:
print(f" ⚠️ {violation['description']}")
cv2.imshow('Mask Detection Test', annotated_frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
frame_count += 1
cap.release()
cv2.destroyAllWindows()
print("\n✅ Test completed!")
print("\nExpected behavior:")
print("- People without masks should show 'VIOLATION' status")
print("- People with masks should show 'COMPLIANT' status")
print("- Masks should be detected with blue bounding boxes")
print("- Violation alerts should appear for missing PPE")
if __name__ == "__main__":
test_violation_logic() |