app.py

from flask import Flask, render_template, request, jsonify
from flask_socketio import SocketIO, emit
from flask_cors import CORS
import cv2
import numpy as np
from ultralytics import YOLO
from scipy.optimize import linear_sum_assignment
from pathlib import Path
import time
import json
from collections import deque
import threading
import base64
import os
import torch


# SAHI imports
try:
    from sahi import AutoDetectionModel
    from sahi.predict import get_sliced_prediction
    from sahi.utils.cv import read_image_as_pil
    SAHI_AVAILABLE = True
    print("✅ SAHI is available")
except ImportError:
    SAHI_AVAILABLE = False
    print("⚠ SAHI not available. Install with: pip install sahi")

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['SECRET_KEY'] = 'secret!'
app.config['MAX_CONTENT_LENGTH'] = 500 * 1024 * 1024  # 500MB max file size

# Allow Next.js frontend (localhost:3000) to call Flask API and Socket.IO
CORS(app, resources={r"/*": {"origins": "*"}})
socketio = SocketIO(
    app,
    async_mode='threading',
    cors_allowed_origins="*",
    ping_timeout=30,
    ping_interval=25,
)

# Ensure upload folder exists
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

class CentroidTracker:
    def __init__(self, max_disappeared=30, max_distance=100):
        self.next_id = 0
        self.tracks = {}
        self.max_disappeared = max_disappeared
        self.max_distance = max_distance

    def register(self, centroid, bbox):
        track_id = self.next_id
        self.tracks[track_id] = {
            'centroid': centroid,
            'bbox': bbox,
            'disappeared': 0
        }
        self.next_id += 1
        return track_id

    def deregister(self, track_id):
        if track_id in self.tracks:
            del self.tracks[track_id]

    def update(self, bboxes):
        if not bboxes:
            for track_id in list(self.tracks.keys()):
                self.tracks[track_id]['disappeared'] += 1
                if self.tracks[track_id]['disappeared'] > self.max_disappeared:
                    self.deregister(track_id)
            return []

        centroids = [(int((x1 + x2) / 2), int((y1 + y2) / 2)) for x1, y1, x2, y2 in bboxes]

        if not self.tracks:
            tracked_objects = [{'id': self.register(c, b), 'bbox': b} for c, b in zip(centroids, bboxes)]
            return tracked_objects

        track_ids = list(self.tracks.keys())
        track_centroids = [self.tracks[tid]['centroid'] for tid in track_ids]
        distances = np.sqrt(((np.array(track_centroids)[:, None] - np.array(centroids)) ** 2).sum(axis=2))

        if distances.size == 0:
            for track_id in track_ids:
                self.tracks[track_id]['disappeared'] += 1
                if self.tracks[track_id]['disappeared'] > self.max_disappeared:
                    self.deregister(track_id)
            return []

        rows, cols = linear_sum_assignment(distances)
        used_rows, used_cols = set(), set()
        tracked_objects = []

        for row, col in zip(rows, cols):
            if distances[row, col] < self.max_distance:
                track_id = track_ids[row]
                self.tracks[track_id].update({
                    'centroid': centroids[col],
                    'bbox': bboxes[col],
                    'disappeared': 0
                })
                tracked_objects.append({'id': track_id, 'bbox': bboxes[col]})
                used_rows.add(row)
                used_cols.add(col)

        for col, (centroid, bbox) in enumerate(zip(centroids, bboxes)):
            if col not in used_cols:
                track_id = self.register(centroid, bbox)
                tracked_objects.append({'id': track_id, 'bbox': bbox})

        for row, track_id in enumerate(track_ids):
            if row not in used_rows:
                self.tracks[track_id]['disappeared'] += 1
                if self.tracks[track_id]['disappeared'] > self.max_disappeared:
                    self.deregister(track_id)

        return tracked_objects

class SafetyMonitor:
    def __init__(self, model_path="best.pt", conf_threshold=0.25, iou_threshold=0.45, 
                 speed_mode="balanced", use_sahi=True, slice_height=640, slice_width=640,
                 overlap_height_ratio=0.2, overlap_width_ratio=0.2):
        """
        Enhanced Safety Monitor with SAHI integration
        """
        self.model_path = model_path
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.speed_mode = speed_mode
        self.use_sahi = use_sahi and SAHI_AVAILABLE
        
        # SAHI parameters
        self.slice_height = slice_height
        self.slice_width = slice_width
        self.overlap_height_ratio = overlap_height_ratio
        self.overlap_width_ratio = overlap_width_ratio
        
        # Initialize models with error handling
        try:
            self.model = YOLO(model_path)
            print(f"✅ Model loaded successfully: {model_path}")
        except Exception as e:
            print(f"❌ Error loading model {model_path}: {e}")
            print("🔄 Trying fallback model...")
            try:
                self.model = YOLO('yolov8n.pt')
                print("✅ Fallback model loaded: yolov8n.pt")
            except Exception as e2:
                print(f"❌ Critical error: Could not load any model: {e2}")
                raise RuntimeError(f"Model loading failed: {e2}")
        
        # Configure device
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"🔧 Using device: {self.device}")
        
        try:
            self.model.to(self.device)
            if self.device == 'cuda' and hasattr(self.model.model, 'half'):
                self.model.model.half()
                print("⚡ Enabled FP16 optimization")
        except Exception as e:
            print(f"⚠ Could not optimize model for {self.device}: {e}")
        
        # Initialize SAHI model
        if self.use_sahi:
            try:
                self.sahi_model = AutoDetectionModel.from_pretrained(
                    model_type='yolov8',
                    model_path=model_path,
                    confidence_threshold=conf_threshold,
                    device=self.device
                )
                print("✅ SAHI model initialized")
            except Exception as e:
                print(f"❌ SAHI initialization failed: {e}")
                self.use_sahi = False
                self.sahi_model = None
        else:
            self.sahi_model = None
            print("ℹ Using regular YOLOv8 inference")
        
        # Speed configuration
        self.speed_config = {
            "fast": {
                "frame_skip": 2,
                "jpeg_quality": 50,
                "max_width": 640,
                "consistency_check_interval": 20,
                "resize_width": 640,
                "resize_height": 480
            },
            "balanced": {
                "frame_skip": 1,
                "jpeg_quality": 60,
                "max_width": 960,
                "consistency_check_interval": 10,
                "resize_width": 800,
                "resize_height": 600
            },
            "accurate": {
                "frame_skip": 0,
                "jpeg_quality": 70,
                "max_width": 1280,
                "consistency_check_interval": 5,
                "resize_width": None,
                "resize_height": None
            }
        }
        
        # Tracking and detection state
        self.person_tracker = CentroidTracker()
        self.expected_classes = ['Helmet', 'Jacket', 'Mask', 'No_Helmet', 'No_Shoes', 
                                'No_goggle', 'No_jacket', 'No_mask', 'Shoes', 'goggle', 'person']
        self.class_names = self.model.names
        self.violations = deque(maxlen=1000)
        self.prev_ppe_detections = {}
        self.ppe_history = {}
        
        # Frame processing optimization
        self.original_dimensions = None
        self.scale_x = 1.0
        self.scale_y = 1.0
        self.frame_counter = 0
        self.current_skip_rate = self.speed_config[speed_mode]["frame_skip"]
        self.processing_times = deque(maxlen=10)
        self.last_processed_frame = None
        self.last_tracked_persons = []
        self.last_ppe_detections = {}
        self.last_violations = []
        self.last_person_compliance = {}
        
        # Thread control
        self.running = False
        self.paused = False
        self.cap = None
        self.thread = None
        self.frame_idx = 0
        self.fps = 0
        self.total_frames = -1
        
        self.load_status_icons()
        self.validate_class_mapping()
        
        # Test model
        try:
            dummy_frame = np.zeros((480, 640, 3), dtype=np.uint8)
            if self.use_sahi:
                self.run_sahi_inference(dummy_frame)
            else:
                self.run_regular_inference(dummy_frame)
            print("✅ Model test successful")
        except Exception as e:
            print(f"❌ Model test failed: {e}")

    def validate_class_mapping(self):
        model_classes = list(self.class_names.values()) if isinstance(self.class_names, dict) else self.class_names
        missing_classes = set(self.expected_classes) - set(model_classes)
        if missing_classes:
            print(f"⚠ Warning: Expected classes not found in model: {missing_classes}")
        extra_classes = set(model_classes) - set(self.expected_classes)
        if extra_classes:
            print(f"ℹ Additional classes found in model: {extra_classes}")

    def load_status_icons(self):
        try:
            icon_size = (30, 30)
            icon_files = {
                'green_helmet': 'static/files/greenHelmet.png',
                'red_helmet': 'static/files/redHelmet.png',
                'green_mask': 'static/files/greenMask.png',
                'red_mask': 'static/files/redMask.png',
                'green_jacket': 'static/files/greenVest.png',
                'red_jacket': 'static/files/redVest.png',
                'green_shoes': 'static/files/greenShoes.png',
                'red_shoes': 'static/files/redShoes.png',
                'green_goggle': 'static/files/greenGoggle.png',
                'red_goggle': 'static/files/redGoggle.png'
            }
            
            for icon_name, file_path in icon_files.items():
                icon = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)
                if icon is not None:
                    icon = cv2.resize(icon, icon_size)
                    setattr(self, icon_name, icon)
                else:
                    print(f"⚠ Warning: Could not load {icon_name} from {file_path}")
                    setattr(self, icon_name, None)
        except Exception as e:
            print(f"❌ Error loading status icons: {e}")

    def resize_frame(self, frame):
        """Resize frame for faster processing"""
        config = self.speed_config[self.speed_mode]
        resize_width = config.get("resize_width")
        resize_height = config.get("resize_height")
        
        if not resize_width or not resize_height:
            return frame
        
        if self.original_dimensions is None:
            self.original_dimensions = (frame.shape[1], frame.shape[0])
            self.scale_x = frame.shape[1] / resize_width
            self.scale_y = frame.shape[0] / resize_height
            print(f"📐 Scale factors: x={self.scale_x:.2f}, y={self.scale_y:.2f}")
        
        resized = cv2.resize(frame, (resize_width, resize_height), interpolation=cv2.INTER_LINEAR)
        return resized
    
    def scale_detections_back(self, boxes):
        """Scale detection boxes back to original frame size"""
        if self.original_dimensions is None or (self.scale_x == 1.0 and self.scale_y == 1.0):
            return boxes
        
        scaled_boxes = []
        for box in boxes:
            if isinstance(box, (list, tuple)):
                x1, y1, x2, y2 = box[:4]
            else:
                x1, y1, x2, y2 = box[:4]
            
            x1_scaled = int(x1 * self.scale_x)
            y1_scaled = int(y1 * self.scale_y)
            x2_scaled = int(x2 * self.scale_x)
            y2_scaled = int(y2 * self.scale_y)
            
            scaled_boxes.append([x1_scaled, y1_scaled, x2_scaled, y2_scaled])
        
        return scaled_boxes

    def should_process_frame(self, frame_idx):
        """Determine if current frame should be processed"""
        if self.current_skip_rate <= 1:
            return True
        return frame_idx % self.current_skip_rate == 0

    def run_sahi_inference(self, frame):
        """Run SAHI inference on frame"""
        try:
            pil_image = read_image_as_pil(frame)
            result = get_sliced_prediction(
                pil_image,
                self.sahi_model,
                slice_height=self.slice_height,
                slice_width=self.slice_width,
                overlap_height_ratio=self.overlap_height_ratio,
                overlap_width_ratio=self.overlap_width_ratio
            )
            
            boxes = []
            confidences = []
            class_names = []
            
            for object_prediction in result.object_prediction_list:
                bbox = object_prediction.bbox
                x1, y1, x2, y2 = bbox.minx, bbox.miny, bbox.maxx, bbox.maxy
                boxes.append([x1, y1, x2, y2])
                confidences.append(object_prediction.score.value)
                class_names.append(object_prediction.category.name)
            
            return boxes, confidences, class_names
            
        except Exception as e:
            print(f"❌ SAHI inference error: {e}")
            return [], [], []

    def run_regular_inference(self, frame):
        """Run regular YOLOv8 inference"""
        try:
            results = self.model(frame, conf=self.conf_threshold, iou=self.iou_threshold, verbose=False)
            result = results[0]
            
            boxes = []
            confidences = []
            class_names = []
            
            if result.boxes is not None:
                boxes = result.boxes.xyxy.cpu().numpy()
                confidences = result.boxes.conf.cpu().numpy()
                class_ids = result.boxes.cls.cpu().numpy().astype(int)
                
                for class_id in class_ids:
                    if class_id < len(self.class_names):
                        class_names.append(self.class_names[class_id])
                    else:
                        class_names.append("unknown")
            
            return boxes, confidences, class_names
            
        except Exception as e:
            print(f"❌ Regular inference error: {e}")
            return [], [], []

    def classify_detection(self, class_name):
        class_name_lower = class_name.lower().strip() if class_name else ''
        
        if class_name_lower == 'person':
            return {'category': 'person', 'is_person': True, 'is_ppe': False, 'is_negative': False}
        
        # Positive PPE detections
        if class_name_lower == 'helmet':
            return {'category': 'ppe_positive', 'is_person': False, 'is_ppe': True, 'is_negative': False, 'ppe_type': 'helmet'}
        if class_name_lower == 'mask':
            return {'category': 'ppe_positive', 'is_person': False, 'is_ppe': True, 'is_negative': False, 'ppe_type': 'mask'}
        if class_name_lower == 'jacket':
            return {'category': 'ppe_positive', 'is_person': False, 'is_ppe': True, 'is_negative': False, 'ppe_type': 'jacket'}
        if class_name_lower == 'shoes':
            return {'category': 'ppe_positive', 'is_person': False, 'is_ppe': True, 'is_negative': False, 'ppe_type': 'shoes'}
        if class_name_lower == 'goggle':
            return {'category': 'ppe_positive', 'is_person': False, 'is_ppe': True, 'is_negative': False, 'ppe_type': 'goggle'}
        
        # Negative PPE detections
        if class_name_lower == 'no_helmet':
            return {'category': 'ppe_negative', 'is_person': False, 'is_ppe': True, 'is_negative': True, 'ppe_type': 'helmet'}
        if class_name_lower == 'no_mask':
            return {'category': 'ppe_negative', 'is_person': False, 'is_ppe': True, 'is_negative': True, 'ppe_type': 'mask'}
        if class_name_lower == 'no_jacket':
            return {'category': 'ppe_negative', 'is_person': False, 'is_ppe': True, 'is_negative': True, 'ppe_type': 'jacket'}
        if class_name_lower == 'no_shoes':
            return {'category': 'ppe_negative', 'is_person': False, 'is_ppe': True, 'is_negative': True, 'ppe_type': 'shoes'}
        if class_name_lower == 'no_goggle':
            return {'category': 'ppe_negative', 'is_person': False, 'is_ppe': True, 'is_negative': True, 'ppe_type': 'goggle'}
        
        return {'category': 'other', 'is_person': False, 'is_ppe': False, 'is_negative': False}

    def check_ppe_compliance(self, person_id, person_bbox, ppe_detections, frame_idx):
        px1, py1, px2, py2 = person_bbox
        compliance = {
            'helmet': False,
            'mask': False,
            'jacket': False,
            'shoes': False,
            'goggle': False,
            'violations': []
        }
        
        if person_id not in self.ppe_history:
            self.ppe_history[person_id] = deque(maxlen=10)
        
        for ppe_type in ['helmet', 'mask', 'jacket', 'shoes', 'goggle']:
            positive_detections = ppe_detections.get(f'{ppe_type}_positive', [])
            negative_detections = ppe_detections.get(f'{ppe_type}_negative', [])
            explicit_violation = False
            
            for ppe_bbox in negative_detections:
                if self.bbox_overlap_check(person_bbox, ppe_bbox, ppe_type):
                    compliance[ppe_type] = False
                    explicit_violation = True
                    break
            
            if not explicit_violation:
                for ppe_bbox in positive_detections:
                    if self.bbox_overlap_check(person_bbox, ppe_bbox, ppe_type):
                        compliance[ppe_type] = True
                        break
            
            self.ppe_history[person_id].append(compliance[ppe_type])
            
            if len(self.ppe_history[person_id]) >= 3:
                recent_states = list(self.ppe_history[person_id])[-3:]
                if all(not state for state in recent_states) and not compliance[ppe_type]:
                    violation_name = f'No {ppe_type.replace("_", " ").title()}'
                    if violation_name not in compliance['violations']:
                        compliance['violations'].append(violation_name)
        
        return compliance

    def bbox_overlap_check(self, person_bbox, ppe_bbox, ppe_type):
        px1, py1, px2, py2 = person_bbox
        ppx1, ppy1, ppx2, ppy2 = ppe_bbox
        ppe_center_x = (ppx1 + ppx2) / 2
        ppe_center_y = (ppy1 + ppy2) / 2
        person_width = px2 - px1
        person_height = py2 - py1
        region_match = False
        
        if ppe_type == 'helmet':
            head_region_bottom = py1 + (person_height * 0.45)
            region_match = (px1 <= ppe_center_x <= px2 and py1 <= ppe_center_y <= head_region_bottom)
        elif ppe_type == 'mask':
            face_region_bottom = py1 + (person_height * 0.55)
            region_match = (px1 <= ppe_center_x <= px2 and py1 <= ppe_center_y <= face_region_bottom)
        elif ppe_type == 'goggle':
            eye_region_bottom = py1 + (person_height * 0.45)
            region_match = (px1 <= ppe_center_x <= px2 and py1 <= ppe_center_y <= eye_region_bottom)
        elif ppe_type == 'jacket':
            torso_top = py1 + (person_height * 0.15)
            torso_bottom = py1 + (person_height * 0.85)
            horizontal_tolerance = person_width * 0.2
            region_match = (px1 - horizontal_tolerance <= ppe_center_x <= px2 + horizontal_tolerance and 
                            torso_top <= ppe_center_y <= torso_bottom)
        elif ppe_type == 'shoes':
            foot_region_top = py1 + (person_height * 0.85)
            horizontal_tolerance = person_width * 0.3
            region_match = (px1 - horizontal_tolerance <= ppe_center_x <= px2 + horizontal_tolerance and 
                            foot_region_top <= ppe_center_y <= py2)
        
        if region_match:
            return True
        
        overlap_x = max(0, min(px2, ppx2) - max(px1, ppx1))
        overlap_y = max(0, min(py2, ppy2) - max(py1, ppy1))
        overlap_area = overlap_x * overlap_y
        
        if overlap_area > 0:
            person_area = (px2 - px1) * (py2 - py1)
            ppe_area = (ppx2 - ppx1) * (ppy2 - ppy1)
            overlap_ratio_person = overlap_area / person_area if person_area > 0 else 0
            overlap_ratio_ppe = overlap_area / ppe_area if ppe_area > 0 else 0
            threshold = 0.15 if ppe_type in ['jacket', 'shoes'] else 0.25
            if overlap_ratio_ppe > threshold or overlap_ratio_person > threshold:
                return True
        
        return False

    def overlay_icon(self, frame, icon, x, y, fallback_text, color):
        if icon is not None and x >= 0 and y >= 0 and x + icon.shape[1] <= frame.shape[1] and y + icon.shape[0] <= frame.shape[0]:
            if len(icon.shape) == 3 and icon.shape[2] == 4:
                alpha = icon[:, :, 3] / 255.0
                for c in range(3):
                    frame[y:y+icon.shape[0], x:x+icon.shape[1], c] = (
                        alpha * icon[:, :, c] + (1 - alpha) * frame[y:y+icon.shape[0], x:x+icon.shape[1], c]
                    )
            else:
                if len(icon.shape) == 3:
                    gray = cv2.cvtColor(icon, cv2.COLOR_BGR2GRAY)
                else:
                    gray = icon
                mask = gray < 240
                frame[y:y+icon.shape[0], x:x+icon.shape[1]][mask] = icon[mask]
        else:
            cv2.putText(frame, fallback_text, (x, y + 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

    def adjust_icon_position(self, frame, x, y, icon_size, occupied_positions):
        icon_w, icon_h = icon_size
        for ox, oy in occupied_positions:
            if abs(x - ox) < icon_w and abs(y - oy) < icon_h:
                x = ox + icon_w + 5
                if x + icon_w > frame.shape[1]:
                    x = ox - icon_w - 5
                if x < 0 or x + icon_w > frame.shape[1]:
                    y += icon_h + 1
                    x = ox
        if x < 0:
            x = 0
        if y < 0:
            y = 0
        if x + icon_w > frame.shape[1]:
            x = frame.shape[1] - icon_w
        if y + icon_h > frame.shape[0]:
            y = frame.shape[0] - icon_h
        return x, y
    

    def check_detection_consistency(self, ppe_detections, frame_idx):
        config = self.speed_config[self.speed_mode]
        if frame_idx % config["consistency_check_interval"] != 0:
            return
            
        for ppe_type in ['helmet', 'mask', 'jacket', 'shoes', 'goggle']:
            pos_key = f'{ppe_type}_positive'
            neg_key = f'{ppe_type}_negative'
            curr_pos = len(ppe_detections.get(pos_key, []))
            curr_neg = len(ppe_detections.get(neg_key, []))
            prev_pos = self.prev_ppe_detections.get(pos_key, 0)
            prev_neg = self.prev_ppe_detections.get(neg_key, 0)
            
            if frame_idx > 0 and (abs(curr_pos - prev_pos) > 3 or abs(curr_neg - prev_neg) > 3):
                print(f"⚠ Detection inconsistency at frame {frame_idx}: {ppe_type} changed from {prev_pos}/{prev_neg} to {curr_pos}/{curr_neg}")
            
            self.prev_ppe_detections[pos_key] = curr_pos
            self.prev_ppe_detections[neg_key] = curr_neg

    def process_frame(self, frame, frame_idx, consistency_interval=10):
        start_time = time.time()
        
        # Check if we should process this frame
        if not self.should_process_frame(frame_idx):
            if self.last_processed_frame is not None:
                return (self.last_processed_frame.copy(), self.last_tracked_persons, 
                       self.last_ppe_detections, self.last_violations, self.last_person_compliance, {
                    'total_persons': len(self.last_tracked_persons),
                    'safe_persons': sum(1 for pc in self.last_person_compliance.values() if not pc.get('violations', [])),
                    'unsafe_persons': len(self.last_tracked_persons) - sum(1 for pc in self.last_person_compliance.values() if not pc.get('violations', [])),
                    'helmet_count': sum(1 for pc in self.last_person_compliance.values() if pc.get('helmet', False)),
                    'jacket_count': sum(1 for pc in self.last_person_compliance.values() if pc.get('jacket', False)),
                    'mask_count': sum(1 for pc in self.last_person_compliance.values() if pc.get('mask', False)),
                    'shoes_count': sum(1 for pc in self.last_person_compliance.values() if pc.get('shoes', False)),
                    'goggle_count': sum(1 for pc in self.last_person_compliance.values() if pc.get('goggle', False)),
                    'fps': self.fps,
                    'processing_time_ms': 0,
                    'progress': (frame_idx / self.total_frames * 100) if self.total_frames > 0 else 0,
                    'skipped': True
                })
        
        # Resize frame for processing
        processed_frame = self.resize_frame(frame)
        
        try:
            # Choose inference method
            if self.use_sahi:
                boxes, confidences, class_names = self.run_sahi_inference(processed_frame)
            else:
                boxes, confidences, class_names = self.run_regular_inference(processed_frame)
            
            # Scale detections back to original size
            if self.original_dimensions is not None:
                boxes = self.scale_detections_back(boxes)
        except Exception as e:
            print(f"❌ Inference error on frame {frame_idx}: {e}")
            annotated_frame = frame.copy()
            cv2.putText(annotated_frame, f"Inference Error: {str(e)[:50]}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            return annotated_frame, [], {}, [], {}, {
                'total_persons': 0, 'safe_persons': 0, 'unsafe_persons': 0,
                'helmet_count': 0, 'jacket_count': 0, 'mask_count': 0, 'shoes_count': 0, 'goggle_count': 0,
                'fps': 0, 'processing_time_ms': 0, 'progress': 0
            }
        
        annotated_frame = frame.copy()
        person_bboxes = []
        
        # Updated PPE detections structure
        ppe_detections = {
            'helmet_positive': [],
            'helmet_negative': [],
            'mask_positive': [],
            'mask_negative': [],
            'jacket_positive': [],
            'jacket_negative': [],
            'shoes_positive': [],
            'shoes_negative': [],
            'goggle_positive': [],
            'goggle_negative': []
        }
        
        # Process detections
        for i, (box, conf, class_name) in enumerate(zip(boxes, confidences, class_names)):
            if conf >= self.conf_threshold:
                if isinstance(box, (list, tuple)):
                    x1, y1, x2, y2 = map(int, box)
                else:
                    x1, y1, x2, y2 = map(int, box[:4])
                
                bbox = [x1, y1, x2, y2]
                classification = self.classify_detection(class_name)
                
                if classification['is_person']:
                    person_bboxes.append(bbox)
                elif classification['is_ppe']:
                    ppe_type = classification['ppe_type']
                    key = f"{ppe_type}_negative" if classification['is_negative'] else f"{ppe_type}_positive"
                    if key in ppe_detections:
                        ppe_detections[key].append(bbox)
        
        self.check_detection_consistency(ppe_detections, frame_idx)
        tracked_persons = self.person_tracker.update(person_bboxes)
        
        person_compliance = {}
        violations = []
        
        for person in tracked_persons:
            person_id = person['id']
            person_bbox = person['bbox']
            compliance = self.check_ppe_compliance(person_id, person_bbox, ppe_detections, frame_idx)
            person_compliance[person_id] = compliance
            
            if compliance['violations']:
                violation = {
                    'frame': frame_idx,
                    'person_id': person_id,
                    'violations': compliance['violations'],
                    'bbox': person_bbox
                }
                violations.append(violation)
                self.violations.append(violation)
        
        # Draw annotations
        if tracked_persons or any(ppe_detections.values()):
            self.draw_annotations(annotated_frame, tracked_persons, person_compliance)
        
        processing_time = time.time() - start_time
        self.fps = 1.0 / processing_time if processing_time > 0 else 0
        
        # Store last results for frame skipping
        self.last_processed_frame = annotated_frame
        self.last_tracked_persons = tracked_persons
        self.last_ppe_detections = ppe_detections
        self.last_violations = violations
        self.last_person_compliance = person_compliance
        
        # Calculate stats
        total_persons = len(tracked_persons)
        safe_persons = sum(1 for pc in person_compliance.values() if not pc['violations'])
        unsafe_persons = total_persons - safe_persons
        
        stats = {
            'total_persons': total_persons,
            'safe_persons': safe_persons,
            'unsafe_persons': unsafe_persons,
            'helmet_count': sum(1 for pc in person_compliance.values() if pc.get('helmet', False)),
            'jacket_count': sum(1 for pc in person_compliance.values() if pc.get('jacket', False)),
            'mask_count': sum(1 for pc in person_compliance.values() if pc.get('mask', False)),
            'shoes_count': sum(1 for pc in person_compliance.values() if pc.get('shoes', False)),
            'goggle_count': sum(1 for pc in person_compliance.values() if pc.get('goggle', False)),
            'fps': round(self.fps, 1),
            'processing_time_ms': round(processing_time * 1000, 1),
            'progress': (frame_idx / self.total_frames * 100) if self.total_frames > 0 else 0,
            'inference_method': 'SAHI' if self.use_sahi else 'Regular',
            'skipped': False
        }
        
        return annotated_frame, tracked_persons, ppe_detections, violations, person_compliance, stats

    def draw_annotations(self, frame, tracked_persons, person_compliance):
        occupied_positions = []
        
        for person in tracked_persons:
            x1, y1, x2, y2 = map(int, person['bbox'])
            track_id = person['id']
            compliance = person_compliance.get(track_id, {})
            violations = compliance.get('violations', [])
            
            color = (0, 255, 0) if not violations else (0, 0, 255)
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3)
            
            label = f"ID-{track_id}"
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)[0]
            cv2.rectangle(frame, (x1, y1 - 35), (x1 + label_size[0] + 10, y1 - 5), color, -1)
            cv2.putText(frame, label, (x1 + 5, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
            
            # PPE status icons
            icon_x_base = x1 - 35
            icon_y_start = y1 - 10
            icon_spacing = 70
            icon_size = (5, 5)
            
            ppe_items = [
                ('helmet', self.green_helmet, self.red_helmet, 'H'),
                ('mask', self.green_mask, self.red_mask, 'M'),
                ('goggle', self.green_goggle, self.red_goggle, 'G'),
                ('jacket', self.green_jacket, self.red_jacket, 'J'),
                ('shoes', self.green_shoes, self.red_shoes, 'S')
            ]
            
            for i, (ppe_type, green_icon, red_icon, text) in enumerate(ppe_items):
                icon_y = icon_y_start + 35 * i
                icon_x, icon_y = self.adjust_icon_position(frame, icon_x_base, icon_y, icon_size, occupied_positions)
                
                is_compliant = compliance.get(ppe_type, False)
                icon_to_use = green_icon if is_compliant else red_icon
                text_color = (0, 255, 0) if is_compliant else (0, 0, 255)
                
                self.overlay_icon(frame, icon_to_use, icon_x, icon_y, text, text_color)
                occupied_positions.append((icon_x, icon_y))

    # def process_video(self, input_path):
    #     try:
    #         if input_path == 0 or str(input_path).lower() == 'webcam':
    #             self.cap = cv2.VideoCapture(1)
    #             config = self.speed_config[self.speed_mode]
    #             self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, config["max_width"])
    #             self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    #             self.total_frames = -1
    #         else:
    #             self.cap = cv2.VideoCapture(input_path)
    #             self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
    #         if not self.cap.isOpened():
    #             raise ValueError(f"Cannot open video source: {input_path}")
            
    #         # Optimize video capture
    #         self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            
    #         config = self.speed_config[self.speed_mode]
    #         frame_skip = config["frame_skip"]
    #         jpeg_quality = config["jpeg_quality"]
    #         max_width = config["max_width"]
    #         consistency_interval = config["consistency_check_interval"]
            
    #         self.frame_idx = 0
    #         last_progress_emit = 0
            
    #         print(f"🚀 Starting video processing in {self.speed_mode} mode:")
    #         print(f"   - Inference: {'SAHI + YOLOv8' if self.use_sahi else 'YOLOv8'}")
    #         print(f"   - Frame skip: {frame_skip}")
    #         print(f"   - JPEG quality: {jpeg_quality}")
    #         print(f"   - Max width: {max_width}px")
            
    #         while self.running and self.cap.isOpened():
    #             if self.paused:
    #                 time.sleep(0.1)
    #                 continue
                
    #             ret, frame = self.cap.read()
    #             if not ret:
    #                 if input_path != 0:
    #                     break
    #                 continue
                
    #             # Resize frame for faster processing
    #             if frame.shape[1] > max_width:
    #                 scale = max_width / frame.shape[1]
    #                 frame = cv2.resize(frame, (0, 0), fx=scale, fy=scale)
                
    #             annotated_frame, tracked_persons, ppe_detections, violations, person_compliance, stats = self.process_frame(frame, self.frame_idx, consistency_interval)
                
    #             # Encode frame with optimized quality
    #             encode_params = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality]
    #             _, buffer = cv2.imencode('.jpg', annotated_frame, encode_params)
    #             frame_base64 = base64.b64encode(buffer).decode('utf-8')
                
    #             # Emit frame data
    #             socketio.emit('frame', {'image': frame_base64})
                
    #             # Emit stats with progress
    #             if self.total_frames > 0:
    #                 progress = min(100, (self.frame_idx / self.total_frames) * 100)
    #                 stats['progress'] = round(progress, 1)
                    
    #                 if progress - last_progress_emit >= 5 or self.frame_idx % 5 == 0:
    #                     socketio.emit('stats', stats)
    #                     last_progress_emit = progress
    #             else:
    #                 if self.frame_idx % 3 == 0:
    #                     socketio.emit('stats', stats)
                
    #             # Emit violations
    #             for violation in violations:
    #                 socketio.emit('violation', violation)
                
    #             self.frame_idx += 1
    #             time.sleep(0.001)
                
    #     except Exception as e:
    #         print(f"❌ Error processing video: {e}")
    #     finally:
    #         if self.cap:
    #             self.cap.release()
    #         self.running = False
    #         socketio.emit('detection_stopped', {'message': 'Video processing completed'})

    def process_video(self, input_path):
        try:
            import depthai as dai
            if input_path == 0 or str(input_path).lower() == 'webcam':
                # Initialize DepthAI pipeline for OAK-D camera
                pipeline = dai.Pipeline()
                cam_rgb = pipeline.create(dai.node.ColorCamera)
                xout_rgb = pipeline.create(dai.node.XLinkOut)
                xout_rgb.setStreamName("rgb")

                # Configure camera properties
                cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
                cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
                config = self.speed_config[self.speed_mode]
                cam_rgb.setVideoSize(config["max_width"], 720)
                cam_rgb.setFps(30)

                # Link camera output to XLinkOut
                cam_rgb.video.link(xout_rgb.input)

                # Start the pipeline
                device = dai.Device(pipeline)
                q_rgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
                self.total_frames = -1
            else:
                # Keep original behavior for file-based input
                self.cap = cv2.VideoCapture(input_path)
                self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
                if not self.cap.isOpened():
                    raise ValueError(f"Cannot open video source: {input_path}")
                self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

            config = self.speed_config[self.speed_mode]
            frame_skip = config["frame_skip"]
            jpeg_quality = config["jpeg_quality"]
            max_width = config["max_width"]
            consistency_interval = config["consistency_check_interval"]

            self.frame_idx = 0
            last_progress_emit = 0

            print(f"🚀 Starting video processing in {self.speed_mode} mode:")
            print(f"   - Inference: {'SAHI + YOLOv8' if self.use_sahi else 'YOLOv8'}")
            print(f"   - Frame skip: {frame_skip}")
            print(f"   - JPEG quality: {jpeg_quality}")
            print(f"   - Max width: {max_width}px")
            print(f"   - Source: {'OAK-D camera' if input_path == 0 or str(input_path).lower() == 'webcam' else 'Video file'}")

            while self.running:
                if self.paused:
                    time.sleep(0.1)
                    continue

                if input_path == 0 or str(input_path).lower() == 'webcam':
                    # Get frame from OAK-D camera
                    in_rgb = q_rgb.tryGet()
                    if in_rgb is None:
                        continue
                    frame = in_rgb.getCvFrame()
                    if frame is None:
                        continue
                else:
                    # Original file-based capture
                    ret, frame = self.cap.read()
                    if not ret:
                        if input_path != 0:
                            break
                        continue

                # Resize frame for faster processing
                if frame.shape[1] > max_width:
                    scale = max_width / frame.shape[1]
                    frame = cv2.resize(frame, (0, 0), fx=scale, fy=scale)

                annotated_frame, tracked_persons, ppe_detections, violations, person_compliance, stats = self.process_frame(frame, self.frame_idx, consistency_interval)

                # Encode frame with optimized quality
                encode_params = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality]
                _, buffer = cv2.imencode('.jpg', annotated_frame, encode_params)
                frame_base64 = base64.b64encode(buffer).decode('utf-8')

                # Emit frame data
                socketio.emit('frame', {'image': frame_base64})

                # Emit stats with progress
                if self.total_frames > 0:
                    progress = min(100, (self.frame_idx / self.total_frames) * 100)
                    stats['progress'] = round(progress, 1)

                    if progress - last_progress_emit >= 5 or self.frame_idx % 5 == 0:
                        socketio.emit('stats', stats)
                        last_progress_emit = progress
                else:
                    if self.frame_idx % 3 == 0:
                        socketio.emit('stats', stats)

                # Emit violations
                for violation in violations:
                    socketio.emit('violation', violation)

                self.frame_idx += 1
                time.sleep(0.001)

        except Exception as e:
            print(f"❌ Error processing video: {e}")
        finally:
            if input_path == 0 or str(input_path).lower() == 'webcam':
                if 'device' in locals():
                    device.close()
            else:
                if self.cap:
                    self.cap.release()
            self.running = False
            socketio.emit('detection_stopped', {'message': 'Video processing completed'})

    def start(self, input_path, conf_threshold, iou_threshold, speed_mode=None):
        if getattr(self, 'running', False):
            try:
                self.stop()
            except Exception as e:
                print(f"⚠ Error stopping previous run before start: {e}")

        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        if speed_mode and speed_mode in self.speed_config:
            self.speed_mode = speed_mode
        self.running = True
        self.paused = False
        self.thread = threading.Thread(target=self.process_video, args=(input_path,))
        self.thread.start()

    def pause(self):
        self.paused = True

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join()
        if self.cap:
            self.cap.release()
        self.cap = None
        self.frame_idx = 0
        self.fps = 0
        self.total_frames = -1

    def start_external(self, conf_threshold, iou_threshold, speed_mode=None):
        if getattr(self, 'running', False):
            try:
                self.stop()
            except Exception as e:
                print(f"⚠ Error stopping previous run before external start: {e}")
        
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        if speed_mode and speed_mode in self.speed_config:
            self.speed_mode = speed_mode
        self.running = True
        self.paused = False
        self.frame_idx = 0
        self.total_frames = -1
        print("✅ External stream started")

    def process_external_frame(self, frame_bgr):
        if not getattr(self, 'running', False):
            return None
        try:
            if frame_bgr is None or frame_bgr.size == 0:
                return None
            
            config = self.speed_config[self.speed_mode]
            jpeg_quality = config["jpeg_quality"]
            max_width = config["max_width"]
            consistency_interval = config["consistency_check_interval"]

            # Resize for speed
            if frame_bgr.shape[1] > max_width:
                scale = max_width / frame_bgr.shape[1]
                frame_bgr = cv2.resize(frame_bgr, (0, 0), fx=scale, fy=scale)

            annotated_frame, tracked_persons, ppe_detections, violations, person_compliance, stats = self.process_frame(
                frame_bgr, self.frame_idx, consistency_interval
            )

            encode_params = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality]
            ok, buffer = cv2.imencode('.jpg', annotated_frame, encode_params)
            if not ok:
                return None
            frame_base64 = base64.b64encode(buffer).decode('utf-8')

            # Emit data
            if self.frame_idx % 2 == 0:
                socketio.emit('frame', {'image': frame_base64})
            socketio.emit('stats', stats)
            for violation in violations:
                socketio.emit('violation', violation)

            self.frame_idx += 1
            return True
        except Exception as e:
            print(f"❌ Error processing external frame: {e}")
            return None

    def stop_external(self):
        self.running = False
        self.paused = False
        self.frame_idx = 0
        self.fps = 0
        self.total_frames = -1
        socketio.emit('detection_stopped', {'message': 'External stream stopped'})
        print("🛑 External stream stopped")

# Initialize monitor with enhanced configuration
try:
    # Try to load custom model first, fallback to default
    try:
        from model_loader import load_ppe_model
        model, model_path = load_ppe_model()
        monitor = SafetyMonitor(
            model_path=model_path,
            use_sahi=SAHI_AVAILABLE,
            speed_mode="balanced"
        )
        print(f"✅ Enhanced PPE Detection system initialized with model: {model_path}")
    except ImportError:
        monitor = SafetyMonitor(
            model_path='best.pt',
            use_sahi=SAHI_AVAILABLE,
            speed_mode="balanced"
        )
        print("✅ Enhanced PPE Detection system initialized with default model")
except Exception as e:
    print(f"❌ Failed to initialize Enhanced PPE Detection system: {e}")
    print("🔄 Falling back to basic detection...")
    monitor = SafetyMonitor(
        model_path='yolov8n.pt',
        use_sahi=False,
        speed_mode="balanced"
    )

# Flask Routes
@app.route('/')
def index():
    return jsonify({"server_status": "running"})

@app.route('/upload', methods=['POST'])
def upload_file():
    try:
        if 'file' not in request.files:
            return jsonify({'success': False, 'error': 'No file part'}), 400
        file = request.files['file']
        if file.filename == '':
            return jsonify({'success': False, 'error': 'No selected file'}), 400
        
        filename = file.filename
        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        abs_path = os.path.abspath(file_path)
        print(f"✅ Uploaded file saved to: {abs_path}")
        return jsonify({'success': True, 'file_path': abs_path})
        
    except Exception as e:
        print(f"❌ Upload error: {e}")
        return jsonify({'success': False, 'error': str(e)}), 500

# Socket.IO Event Handlers
@socketio.on('start_detection')
def handle_start_detection(data):
    input_path = data.get('input_path', 0)
    conf = float(data.get('conf', 0.25))
    iou = float(data.get('iou', 0.45))
    speed_mode = data.get('speed_mode', 'balanced')
    use_sahi = data.get('use_sahi', True)
    
    # Update SAHI setting if available
    if SAHI_AVAILABLE:
        monitor.use_sahi = use_sahi
    
    monitor.start(input_path, conf, iou, speed_mode)

@socketio.on('connect')
def handle_connect():
    try:
        print(f"✅ Socket.IO client connected")
        emit('server_info', {
            'message': 'Connected to Enhanced Flask-SocketIO PPE Monitor',
            'sahi_available': SAHI_AVAILABLE,
            'current_mode': monitor.speed_mode,
            'using_sahi': monitor.use_sahi
        })
    except Exception as e:
        print(f"Socket connect handler error: {e}")

@socketio.on('disconnect')
def handle_disconnect():
    print('⚠ Socket.IO client disconnected')

@socketio.on('pause_detection')
def handle_pause_detection():
    monitor.pause()

@socketio.on('stop_detection')
def handle_stop_detection():
    monitor.stop()

@socketio.on('update_settings')
def handle_update_settings(data):
    monitor.conf_threshold = float(data.get('conf', monitor.conf_threshold))
    monitor.iou_threshold = float(data.get('iou', monitor.iou_threshold))
    speed_mode = data.get('speed_mode')
    if speed_mode and speed_mode in monitor.speed_config:
        monitor.speed_mode = speed_mode
    use_sahi = data.get('use_sahi')
    if use_sahi is not None and SAHI_AVAILABLE:
        monitor.use_sahi = use_sahi

@socketio.on('toggle_sahi')
def handle_toggle_sahi():
    if SAHI_AVAILABLE:
        monitor.use_sahi = not monitor.use_sahi
        emit('sahi_toggled', {'using_sahi': monitor.use_sahi})
        print(f"🔄 SAHI toggled: {'ON' if monitor.use_sahi else 'OFF'}")
    else:
        emit('sahi_error', {'message': 'SAHI not available'})

# REST API Endpoints for external streaming
@app.route('/stream/start', methods=['POST'])
def stream_start():
    try:
        data = request.get_json(silent=True) or {}
        conf = float(data.get('conf', 0.25))
        iou = float(data.get('iou', 0.45))
        speed_mode = data.get('speed_mode', 'balanced')
        use_sahi = data.get('use_sahi', True)
        
        if SAHI_AVAILABLE:
            monitor.use_sahi = use_sahi
            
        monitor.start_external(conf, iou, speed_mode)
        return jsonify({
            'success': True,
            'sahi_available': SAHI_AVAILABLE,
            'using_sahi': monitor.use_sahi,
            'speed_mode': speed_mode
        })
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

@app.route('/stream/frame', methods=['POST'])
def stream_frame():
    try:
        payload = request.get_json(silent=True)
        if not payload or 'image' not in payload:
            return jsonify({'success': False, 'error': 'No image provided'}), 400
        
        image_data = payload['image']
        if isinstance(image_data, str) and image_data.startswith('data:image'):
            header, b64data = image_data.split(',', 1)
        else:
            b64data = image_data
        
        if not b64data:
            return jsonify({'success': False, 'error': 'Empty image payload'}), 400
        
        try:
            img_bytes = base64.b64decode(b64data)
        except Exception:
            return jsonify({'success': False, 'error': 'Invalid base64 data'}), 400
        
        np_arr = np.frombuffer(img_bytes, np.uint8)
        if np_arr.size == 0:
            return jsonify({'success': False, 'error': 'Empty frame data'}), 400
        
        frame_bgr = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        if frame_bgr is None:
            return jsonify({'success': False, 'error': 'Invalid image data'}), 400
        
        processed = monitor.process_external_frame(frame_bgr)
        return jsonify({'success': bool(processed)})
        
    except Exception as e:
        print(f"❌ stream_frame error: {e}")
        return jsonify({'success': False, 'error': str(e)}), 500

@app.route('/stream/stop', methods=['POST'])
def stream_stop():
    try:
        monitor.stop_external()
        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

@app.route('/api/status', methods=['GET'])
def get_status():
    return jsonify({
        'running': monitor.running,
        'paused': monitor.paused,
        'frame_idx': monitor.frame_idx,
        'fps': monitor.fps,
        'total_frames': monitor.total_frames,
        'sahi_available': SAHI_AVAILABLE,
        'using_sahi': monitor.use_sahi,
        'speed_mode': monitor.speed_mode,
        'conf_threshold': monitor.conf_threshold,
        'iou_threshold': monitor.iou_threshold
    })

if __name__ == '__main__':
    print("🚀 Starting Enhanced PPE Detection Flask Server...")
    print(f"📊 SAHI Support: {'✅ Available' if SAHI_AVAILABLE else '❌ Not Available'}")
    print(f"🎯 Default Mode: {monitor.speed_mode}")
    print(f"🔍 Inference Method: {'SAHI + YOLOv8' if monitor.use_sahi else 'YOLOv8'}")
    socketio.run(app, debug=True, host='0.0.0.0', port=5500, allow_unsafe_werkzeug=True)