Upload safety_detector.py

safety_detector.py

@@ -0,0 +1,926 @@
+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()