import cv2 import numpy as np from ultralytics import YOLO import torch import time from datetime import datetime import os import json from threading import Thread import queue from typing import Dict, List, Tuple, Optional import requests class SafetyDetector: """ Real-time safety compliance detection system using YOLO for object detection. Detects people and safety equipment like hard hats, safety vests, and safety glasses. """ def __init__(self, model_path: Optional[str] = None, confidence_threshold: float = 0.5): """ Initialize the safety detector with a specialized PPE detection model. Args: model_path: Path to custom model, if None will download PPE model confidence_threshold: Minimum confidence for detections """ self.confidence_threshold = confidence_threshold self.device = 'cuda' if torch.cuda.is_available() else 'cpu' # Stricter confidence thresholds for different equipment types to reduce false positives self.equipment_confidence_thresholds = { 'hardhat': 0.7, # Higher threshold for hard hats (hair confusion) 'safety_vest': 0.75, # Higher threshold for safety vests (clothing confusion) 'mask': 0.6, # Moderate threshold for masks 'person': 0.5, # Standard threshold for people 'no_hardhat': 0.6, # Moderate threshold for NO- detections 'no_safety_vest': 0.6, 'no_mask': 0.6 } # Try to load a specialized PPE detection model self.model = self._load_ppe_model(model_path) # PPE class names - these are the actual classes we expect from PPE models self.ppe_classes = { 'hardhat': ['Hardhat', 'hardhat', 'helmet', 'hard hat'], 'safety_vest': ['Safety Vest', 'safety vest', 'vest', 'safety-vest', 'Safety-Vest'], 'no_hardhat': ['NO-Hardhat', 'no-hardhat', 'no hardhat', 'NO-Helmet'], 'no_safety_vest': ['NO-Safety Vest', 'no-safety-vest', 'no safety vest', 'NO-Safety-Vest'], 'person': ['Person', 'person'], 'mask': ['Mask', 'mask'], 'no_mask': ['NO-Mask', 'no-mask', 'no mask'], 'safety_gloves': ['Safety Gloves', 'safety-gloves', 'gloves', 'Gloves'], 'safety_glasses': ['Safety Glasses', 'safety-glasses', 'glasses', 'Safety-Glasses'], 'hearing_protection': ['Hearing Protection', 'hearing-protection', 'ear protection'] } print(f"Using device: {self.device}") print(f"Loaded PPE detection model with stricter confidence thresholds") print(f"Equipment thresholds: {self.equipment_confidence_thresholds}") # Colors for bounding boxes self.colors = { 'person': (0, 255, 0), # Green for compliant person 'violation': (0, 0, 255), # Red for safety violation 'equipment': (255, 255, 0), # Yellow for safety equipment 'warning': (0, 165, 255) # Orange for warnings } # Violation tracking self.violations = [] self.violation_images_dir = "violation_captures" os.makedirs(self.violation_images_dir, exist_ok=True) def _load_ppe_model(self, model_path: Optional[str] = None) -> YOLO: """Load a specialized PPE detection model.""" if model_path and os.path.exists(model_path): print(f"Loading custom model from {model_path}") return YOLO(model_path) # Try to download YOLOv8-compatible PPE models ppe_model_urls = [ # Try the snehilsanyal YOLOv8 PPE model (best.pt) "https://github.com/snehilsanyal/Construction-Site-Safety-PPE-Detection/raw/main/models/best.pt", # Try mayank13-01 YOLOv8 PPE model "https://github.com/mayank13-01/Yolov8-PPE/raw/main/YOLO-Weights/ppe.pt" ] for i, url in enumerate(ppe_model_urls): try: model_filename = f"ppe_yolov8_model_{i}.pt" if not os.path.exists(model_filename): print(f"Downloading PPE detection model from {url}...") response = requests.get(url, timeout=60) if response.status_code == 200: with open(model_filename, 'wb') as f: f.write(response.content) print(f"Downloaded PPE model successfully as {model_filename}") if os.path.exists(model_filename): print(f"Loading YOLOv8 PPE model from {model_filename}") model = YOLO(model_filename) # Test if the model loads properly classes = self._get_model_classes(model) print(f"Model classes: {classes}") # Check if it has PPE-related classes ppe_related = any( any(keyword in str(cls).lower() for keyword in ['hardhat', 'vest', 'helmet', 'mask', 'person']) for cls in classes ) if ppe_related: print(f"✅ Found PPE-capable model with {len(classes)} classes") return model else: print(f"⚠️ Model doesn't seem to have PPE classes: {classes}") except Exception as e: print(f"Failed to download/load from {url}: {e}") continue # Fallback to YOLOv8 with a warning print("⚠️ Warning: Could not load specialized PPE model, falling back to YOLOv8n") print(" Note: YOLOv8n can detect people but not safety equipment") return YOLO('yolov8n.pt') def _get_model_classes(self, model=None) -> List[str]: """Get the list of classes the model can detect.""" if model is None: model = self.model if hasattr(model, 'names'): return list(model.names.values()) return [] def _get_class_category(self, class_name: str) -> str: """Map detected class name to our safety categories.""" class_name_lower = class_name.lower() for category, variations in self.ppe_classes.items(): for variation in variations: if variation.lower() in class_name_lower or class_name_lower in variation.lower(): return category return class_name_lower def detect_safety_violations(self, frame: np.ndarray) -> Dict: """ Detect safety violations in the given frame with improved accuracy. Returns: Dictionary containing detection results and violations """ start_time = time.time() # Run detection with optimized settings for speed results = self.model(frame, conf=0.3, verbose=False, imgsz=640, half=False) detections = [] people_count = 0 safety_equipment_detected = { 'hardhat': 0, 'safety_vest': 0, 'safety_gloves': 0, 'safety_glasses': 0, 'hearing_protection': 0, 'mask': 0 } violations = [] no_equipment_detections = [] # Track NO- detections separately # Process detections with stricter filtering for r in results: boxes = r.boxes if boxes is not None: for box in boxes: # Get detection info x1, y1, x2, y2 = box.xyxy[0].cpu().numpy() confidence = box.conf[0].cpu().numpy() class_id = int(box.cls[0].cpu().numpy()) # Get class name if hasattr(self.model, 'names'): class_name = self.model.names[class_id] else: class_name = f"class_{class_id}" # Map to our categories category = self._get_class_category(class_name) # Apply stricter confidence thresholds based on equipment type required_confidence = self.equipment_confidence_thresholds.get(category, self.confidence_threshold) # Skip detections that don't meet the stricter threshold if confidence < required_confidence: continue detection = { 'bbox': [int(x1), int(y1), int(x2), int(y2)], 'confidence': float(confidence), 'class': class_name, 'category': category } detections.append(detection) # Count people and safety equipment if category == 'person': people_count += 1 elif category in safety_equipment_detected: safety_equipment_detected[category] += 1 elif category in ['hardhat', 'safety_vest', 'mask'] and not category.startswith('no_'): safety_equipment_detected[category] += 1 # Handle negative detections (NO-Hardhat, NO-Mask, etc.) # These indicate violations - a person without required equipment if category.startswith('no_'): equipment_type = category.replace('no_', '') if equipment_type in ['hardhat', 'safety_vest', 'mask']: no_equipment_detections.append({ 'type': f'missing_{equipment_type}', 'severity': 'high', 'description': f'Person detected without {equipment_type.replace("_", " ").title()}', 'bbox': [int(x1), int(y1), int(x2), int(y2)], 'confidence': float(confidence), 'equipment_type': equipment_type }) # Create violations based on NO- detections (these are more reliable) violations.extend(no_equipment_detections) # If we have people but no NO- detections, check equipment ratios if people_count > 0 and len(no_equipment_detections) == 0: required_equipment = ['hardhat', 'safety_vest', 'mask'] for equipment in required_equipment: detected_count = safety_equipment_detected[equipment] # If significantly fewer equipment than people, assume violations if detected_count < people_count * 0.8: # Allow some tolerance missing_count = people_count - detected_count equipment_name = equipment.replace("_", " ").title() violations.append({ 'type': f'missing_{equipment}', 'severity': 'high', 'description': f'{missing_count} person(s) likely missing {equipment_name}', 'count': missing_count }) # Special handling for masks - they're often not detected well mask_detected = safety_equipment_detected['mask'] no_mask_detected = len([v for v in no_equipment_detections if v['equipment_type'] == 'mask']) if people_count > 0 and mask_detected == 0 and no_mask_detected == 0: # No mask detections at all - assume people are not wearing masks violations.append({ 'type': 'missing_mask', 'severity': 'high', 'description': f'{people_count} person(s) not wearing Face Mask', 'count': people_count }) processing_time = time.time() - start_time return { 'detections': detections, 'people_count': people_count, 'safety_equipment': safety_equipment_detected, 'violations': violations, 'processing_time': processing_time, 'fps': 1.0 / processing_time if processing_time > 0 else 0 } def draw_detections(self, frame: np.ndarray, results: Dict) -> np.ndarray: """ Draw premium bounding boxes only for POSITIVE equipment detections. No boxes for missing equipment - violations shown through person status only. Args: frame: Input frame results: Detection results containing detections, violations, etc. Returns: Annotated frame with premium styling """ annotated_frame = frame.copy() height, width = annotated_frame.shape[:2] # Create overlay for semi-transparent effects overlay = annotated_frame.copy() # Premium color scheme colors = { 'person_compliant': (46, 204, 113), # Emerald green 'person_violation': (231, 76, 60), # Red 'equipment': (52, 152, 219), # Blue 'hardhat': (46, 204, 113), # Green 'safety_vest': (241, 196, 15), # Yellow 'mask': (0, 191, 255), # Deep sky blue 'violation_bg': (231, 76, 60), # Red background 'text_bg': (44, 62, 80), # Dark blue-gray 'text_primary': (255, 255, 255), # White 'text_secondary': (149, 165, 166), # Light gray 'shadow': (0, 0, 0), # Black shadow 'accent': (155, 89, 182), # Purple accent } # Track people and their compliance status people_status = {} # First pass: categorize people for detection in results.get('detections', []): class_name = detection['class'].lower() bbox = detection['bbox'] confidence = detection['confidence'] if 'person' in class_name: person_id = f"person_{bbox[0]}_{bbox[1]}" people_status[person_id] = { 'bbox': bbox, 'confidence': confidence, 'violations': [], 'equipment': [] } # Map violations to people for violation in results.get('violations', []): if 'bbox' in violation: # This is a specific violation with a bounding box (from NO- detections) violation_bbox = violation['bbox'] # Find the closest person to this violation closest_person = None min_distance = float('inf') for person_id, person_data in people_status.items(): person_bbox = person_data['bbox'] # Calculate distance between violation and person distance = abs(violation_bbox[0] - person_bbox[0]) + abs(violation_bbox[1] - person_bbox[1]) if distance < min_distance: min_distance = distance closest_person = person_id if closest_person and min_distance < 100: # Within reasonable distance violation_type = violation['type'].replace('missing_', '') people_status[closest_person]['violations'].append(violation_type) else: # General violation - apply to all people (when equipment count < people count) violation_type = violation['type'].replace('missing_', '') for person_id in people_status: people_status[person_id]['violations'].append(violation_type) # If no specific violations detected but people are present, assume they're missing all required equipment if len(people_status) > 0 and len(results.get('violations', [])) == 0: # Check if we have any positive equipment detections equipment_detected = any( detection['category'] in ['hardhat', 'safety_vest', 'mask'] for detection in results.get('detections', []) if detection['category'] in ['hardhat', 'safety_vest', 'mask'] ) # If no equipment detected at all, mark all people as having violations if not equipment_detected: for person_id in people_status: people_status[person_id]['violations'] = ['hardhat', 'safety_vest', 'mask'] # ONLY draw POSITIVE equipment detections (when equipment IS being worn) for detection in results.get('detections', []): class_name = detection['class'].lower() category = detection.get('category', '') # Skip people and NO- detections - we only want positive equipment if 'person' in class_name or 'no-' in class_name or 'no_' in category: continue # Only draw positive equipment detections if category in ['hardhat', 'safety_vest', 'mask'] or any(equip in class_name for equip in ['hardhat', 'vest', 'helmet', 'safety', 'mask']): bbox = detection['bbox'] confidence = detection['confidence'] # Choose color and label based on equipment type if any(x in class_name for x in ['hardhat', 'helmet']) or category == 'hardhat': color = colors['hardhat'] equipment_type = "Hard Hat ✓" elif 'vest' in class_name or category == 'safety_vest': color = colors['safety_vest'] equipment_type = "Safety Vest ✓" elif 'mask' in class_name or category == 'mask': color = colors['mask'] equipment_type = "Face Mask ✓" else: color = colors['equipment'] equipment_type = "Safety Equipment ✓" # Draw equipment with premium styling self._draw_premium_bbox(overlay, annotated_frame, bbox, color, equipment_type, confidence, bbox_type="equipment", colors=colors) # Draw people with compliance status (no violation indicators on person boxes) for person_id, person_data in people_status.items(): bbox = person_data['bbox'] confidence = person_data['confidence'] violations = person_data['violations'] # Determine person status is_compliant = len(violations) == 0 color = colors['person_compliant'] if is_compliant else colors['person_violation'] status_text = "COMPLIANT" if is_compliant else "VIOLATION" # Draw person with premium styling (no violation details on the box) self._draw_premium_bbox(overlay, annotated_frame, bbox, color, f"Person - {status_text}", confidence, bbox_type="person", violations=None, # Don't show violation details on person box colors=colors) # Blend overlay with original frame for semi-transparent effects alpha = 0.15 cv2.addWeighted(overlay, alpha, annotated_frame, 1 - alpha, 0, annotated_frame) # Statistics are now handled by the web UI, no overlay needed on video feed return annotated_frame def _draw_premium_bbox(self, overlay, frame, bbox, color, label, confidence, bbox_type="default", violations=None, colors=None): """Draw a premium-styled bounding box with advanced visual effects.""" x1, y1, x2, y2 = map(int, bbox) # Box dimensions box_width = x2 - x1 box_height = y2 - y1 # Draw shadow first (slightly offset) shadow_offset = 3 shadow_color = colors['shadow'] cv2.rectangle(overlay, (x1 + shadow_offset, y1 + shadow_offset), (x2 + shadow_offset, y2 + shadow_offset), shadow_color, 2) # Main bounding box with thinner lines box_thickness = 2 if bbox_type == "person" else 1 # Draw main rectangle cv2.rectangle(frame, (x1, y1), (x2, y2), color, box_thickness) # Draw corner accents for premium look corner_length = min(20, box_width // 4, box_height // 4) accent_thickness = box_thickness # Top-left corner cv2.line(frame, (x1, y1), (x1 + corner_length, y1), color, accent_thickness) cv2.line(frame, (x1, y1), (x1, y1 + corner_length), color, accent_thickness) # Top-right corner cv2.line(frame, (x2, y1), (x2 - corner_length, y1), color, accent_thickness) cv2.line(frame, (x2, y1), (x2, y1 + corner_length), color, accent_thickness) # Bottom-left corner cv2.line(frame, (x1, y2), (x1 + corner_length, y2), color, accent_thickness) cv2.line(frame, (x1, y2), (x1, y2 - corner_length), color, accent_thickness) # Bottom-right corner cv2.line(frame, (x2, y2), (x2 - corner_length, y2), color, accent_thickness) cv2.line(frame, (x2, y2), (x2, y2 - corner_length), color, accent_thickness) # Prepare label text confidence_text = f"{confidence:.1%}" main_text = f"{label}" # Calculate text dimensions font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.5 thickness = 1 (main_w, main_h), _ = cv2.getTextSize(main_text, font, font_scale, thickness) (conf_w, conf_h), _ = cv2.getTextSize(confidence_text, font, font_scale - 0.1, thickness - 1) # Label background dimensions label_height = max(main_h, conf_h) + 12 label_width = max(main_w, conf_w) + 16 # Position label (above box if space available, otherwise below) if y1 - label_height - 5 > 0: label_y = y1 - label_height - 5 else: label_y = y2 + 5 label_x = x1 # Ensure label stays within frame if label_x + label_width > frame.shape[1]: label_x = frame.shape[1] - label_width - 5 if label_x < 0: label_x = 5 # Draw label background with gradient effect bg_color = colors['text_bg'] # Main background cv2.rectangle(overlay, (label_x, label_y), (label_x + label_width, label_y + label_height), bg_color, -1) # Colored top border cv2.rectangle(frame, (label_x, label_y), (label_x + label_width, label_y + 4), color, -1) # Add subtle border cv2.rectangle(frame, (label_x, label_y), (label_x + label_width, label_y + label_height), color, 1) # Draw main text text_y = label_y + main_h + 6 cv2.putText(frame, main_text, (label_x + 8, text_y), font, font_scale, colors['text_primary'], thickness) # Draw confidence text conf_y = text_y + conf_h + 4 cv2.putText(frame, confidence_text, (label_x + 8, conf_y), font, font_scale - 0.1, colors['text_secondary'], max(1, thickness - 1)) # Draw violation indicators for people (only if violations are provided) if bbox_type == "person" and violations is not None and len(violations) > 0: self._draw_violation_indicators(frame, overlay, x1, y1, x2, y2, violations, colors) def _draw_violation_indicators(self, frame, overlay, x1, y1, x2, y2, violations, colors): """Draw violation indicators with premium styling.""" # Warning icon position (top-right of bounding box) icon_size = 24 icon_x = x2 - icon_size - 5 icon_y = y1 + 5 # Draw warning background circle cv2.circle(overlay, (icon_x + icon_size//2, icon_y + icon_size//2), icon_size//2, colors['violation_bg'], -1) cv2.circle(frame, (icon_x + icon_size//2, icon_y + icon_size//2), icon_size//2, colors['violation_bg'], 2) # Draw exclamation mark center_x = icon_x + icon_size//2 center_y = icon_y + icon_size//2 # Exclamation line cv2.line(frame, (center_x, center_y - 6), (center_x, center_y + 2), colors['text_primary'], 2) # Exclamation dot cv2.circle(frame, (center_x, center_y + 5), 1, colors['text_primary'], -1) # Draw violation list below the person if space allows violation_text = "Missing: " + ", ".join(violations) font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.5 thickness = 1 (text_w, text_h), _ = cv2.getTextSize(violation_text, font, font_scale, thickness) # Position violation text viol_x = x1 viol_y = y2 + text_h + 8 # Ensure text stays within frame if viol_y + text_h > frame.shape[0]: viol_y = y1 - text_h - 8 if viol_x + text_w > frame.shape[1]: viol_x = frame.shape[1] - text_w - 5 # Draw violation text background padding = 4 cv2.rectangle(overlay, (viol_x - padding, viol_y - text_h - padding), (viol_x + text_w + padding, viol_y + padding), colors['violation_bg'], -1) # Draw violation text cv2.putText(frame, violation_text, (viol_x, viol_y), font, font_scale, colors['text_primary'], thickness) def _draw_statistics_overlay(self, frame, results, colors, width, height): """Draw statistics overlay with premium styling.""" # Statistics data people_count = results.get('people_count', 0) violations = results.get('violations', []) violation_count = len(violations) compliant_count = people_count - violation_count compliance_rate = (compliant_count / max(people_count, 1)) * 100 # Statistics text stats = [ f"People: {people_count}", f"Compliant: {compliant_count}", f"Violations: {violation_count}", f"Compliance: {compliance_rate:.1f}%" ] # Text properties font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.7 thickness = 2 # Calculate background size max_text_width = 0 total_height = 0 line_heights = [] for text in stats: (text_w, text_h), _ = cv2.getTextSize(text, font, font_scale, thickness) max_text_width = max(max_text_width, text_w) line_heights.append(text_h) total_height += text_h + 8 # Background dimensions bg_width = max_text_width + 24 bg_height = total_height + 16 # Position (top-left corner) bg_x = 20 bg_y = 20 # Draw semi-transparent background overlay = frame.copy() cv2.rectangle(overlay, (bg_x, bg_y), (bg_x + bg_width, bg_y + bg_height), colors['text_bg'], -1) cv2.addWeighted(overlay, 0.8, frame, 0.2, 0, frame) # Draw border cv2.rectangle(frame, (bg_x, bg_y), (bg_x + bg_width, bg_y + bg_height), colors['accent'], 2) # Draw statistics text current_y = bg_y + 24 for i, text in enumerate(stats): # Choose color based on statistic type if "Violations:" in text and violation_count > 0: text_color = colors['person_violation'] elif "Compliant:" in text: text_color = colors['person_compliant'] elif "Compliance:" in text: if compliance_rate >= 80: text_color = colors['person_compliant'] elif compliance_rate >= 60: text_color = colors['safety_vest'] else: text_color = colors['person_violation'] else: text_color = colors['text_primary'] cv2.putText(frame, text, (bg_x + 12, current_y), font, font_scale, text_color, thickness) current_y += line_heights[i] + 8 def get_model_classes(self) -> List[str]: """Get the list of classes the model can detect.""" return self._get_model_classes() def test_detection(self, test_image_path: str = None): """Test the detector with a sample image or webcam.""" if test_image_path and os.path.exists(test_image_path): frame = cv2.imread(test_image_path) if frame is not None: results = self.detect_safety_violations(frame) output = self.draw_detections(frame, results) print(f"Detected classes: {[d['class'] for d in results['detections']]}") print(f"Available model classes: {self.get_model_classes()}") cv2.imshow('PPE Detection Test', output) cv2.waitKey(0) cv2.destroyAllWindows() return results else: print("Testing with webcam - press 'q' to quit") cap = cv2.VideoCapture(0) while True: ret, frame = cap.read() if not ret: break results = self.detect_safety_violations(frame) output = self.draw_detections(frame, results) cv2.imshow('PPE Detection Test', output) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() cv2.destroyAllWindows() def analyze_safety_compliance(self, detections: List[Dict]) -> Dict: """ Analyze safety compliance based on detected objects. Args: detections: List of detected objects Returns: Dictionary with compliance analysis """ people_detected = [] safety_equipment = [] # Separate people and safety equipment for detection in detections: if detection['class'].lower() == 'person': people_detected.append(detection) elif any(equipment in detection['class'].lower() for equipment in ['helmet', 'hardhat', 'vest', 'gloves', 'glasses']): safety_equipment.append(detection) # Analyze compliance for each person compliance_results = [] for person in people_detected: person_bbox = person['bbox'] # Check for nearby safety equipment nearby_equipment = self._find_nearby_equipment(person_bbox, safety_equipment) # Determine missing equipment required_equipment = ['hardhat', 'safety_vest'] missing_equipment = [] for equipment in required_equipment: if not any(equipment.lower() in item['class'].lower() for item in nearby_equipment): missing_equipment.append(equipment) compliance_results.append({ 'person': person, 'nearby_equipment': nearby_equipment, 'missing_equipment': missing_equipment, 'is_compliant': len(missing_equipment) == 0, 'compliance_score': 1.0 - (len(missing_equipment) / len(required_equipment)) }) return { 'total_people': len(people_detected), 'compliant_people': sum(1 for result in compliance_results if result['is_compliant']), 'violations': sum(len(result['missing_equipment']) for result in compliance_results), 'compliance_results': compliance_results, 'overall_compliance_rate': ( sum(result['compliance_score'] for result in compliance_results) / max(len(compliance_results), 1) ) } def _find_nearby_equipment(self, person_bbox: List[int], equipment_list: List[Dict], proximity_threshold: float = 0.3) -> List[Dict]: """Find safety equipment near a person.""" nearby_equipment = [] person_center_x = (person_bbox[0] + person_bbox[2]) / 2 person_center_y = (person_bbox[1] + person_bbox[3]) / 2 for equipment in equipment_list: equip_bbox = equipment['bbox'] equip_center_x = (equip_bbox[0] + equip_bbox[2]) / 2 equip_center_y = (equip_bbox[1] + equip_bbox[3]) / 2 # Calculate normalized distance distance = np.sqrt((person_center_x - equip_center_x)**2 + (person_center_y - equip_center_y)**2) # Normalize by image diagonal (assuming standard frame size) normalized_distance = distance / 1000 # Adjust based on typical frame size if normalized_distance < proximity_threshold: nearby_equipment.append(equipment) return nearby_equipment def draw_annotations(self, frame: np.ndarray, analysis: Dict) -> np.ndarray: """ Draw bounding boxes and annotations on the frame. Args: frame: Input frame analysis: Safety compliance analysis results Returns: Annotated frame """ annotated_frame = frame.copy() # Draw safety equipment for equipment in analysis['safety_equipment']: bbox = equipment['bbox'] cv2.rectangle(annotated_frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), self.colors['equipment'], 2) label = f"{equipment.get('equipment_type', equipment['class'])}: {equipment['confidence']:.2f}" cv2.putText(annotated_frame, label, (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors['equipment'], 2) # Draw people with compliance status for result in analysis['compliance_results']: person = result['person'] bbox = person['bbox'] # Choose color based on compliance color = self.colors['person'] if result['is_compliant'] else self.colors['violation'] # Draw bounding box cv2.rectangle(annotated_frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 3) # Create status label status = "COMPLIANT" if result['is_compliant'] else "VIOLATION" confidence_text = f"Person: {person['confidence']:.2f}" # Draw labels cv2.putText(annotated_frame, status, (bbox[0], bbox[1] - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2) cv2.putText(annotated_frame, confidence_text, (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) # Show missing equipment if result['missing_equipment']: missing_text = f"Missing: {', '.join(result['missing_equipment'])}" cv2.putText(annotated_frame, missing_text, (bbox[0], bbox[3] + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors['violation'], 2) # Draw summary statistics summary_text = [ f"Total People: {analysis['total_people']}", f"Compliant: {analysis['compliant_people']}", f"Violations: {analysis['violations']}", f"Compliance Rate: {(analysis['compliant_people']/max(analysis['total_people'],1)*100):.1f}%" ] for i, text in enumerate(summary_text): cv2.putText(annotated_frame, text, (10, 30 + i * 25), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) return annotated_frame def capture_violation(self, frame: np.ndarray, violation_data: Dict) -> str: """ Capture and save an image when a safety violation is detected. Args: frame: Current frame violation_data: Information about the violation Returns: Path to saved image """ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3] filename = f"violation_{timestamp}.jpg" filepath = os.path.join(self.violation_images_dir, filename) # Save the frame cv2.imwrite(filepath, frame) # Save violation metadata metadata = { 'timestamp': datetime.now().isoformat(), 'filename': filename, 'violation_data': violation_data } metadata_file = filepath.replace('.jpg', '_metadata.json') with open(metadata_file, 'w') as f: json.dump(metadata, f, indent=2) self.violations.append(metadata) return filepath def process_frame(self, frame: np.ndarray) -> Tuple[np.ndarray, Dict]: """ Process a single frame for safety monitoring. Args: frame: Input video frame Returns: Tuple of (annotated_frame, analysis_results) """ # Detect objects and get safety violations results = self.detect_safety_violations(frame) # Draw detections on frame using the main drawing method annotated_frame = self.draw_detections(frame, results) return annotated_frame, { 'detections': results['detections'], 'people_count': results['people_count'], 'safety_equipment': results['safety_equipment'], 'violations': results['violations'], 'violation_summary': self.get_violation_summary(), 'frame_stats': { 'processing_time': results['processing_time'], 'fps': results['fps'], 'detection_count': len(results['detections']) } } def get_violation_summary(self) -> Dict: """Get a summary of recent violations.""" # This would typically connect to a database or log file # For now, return a placeholder return { 'total_violations_today': 0, 'most_common_violation': 'missing_hardhat', 'compliance_trend': [] # Could track compliance over time } if __name__ == "__main__": # Test the detector detector = SafetyDetector() print("Available classes:", detector.get_model_classes()) detector.test_detection()