import cv2
from cvzone.FaceMeshModule import FaceMeshDetector

class RailwayFatigueDetector:
    """
    铁路驾驶员疲劳检测器
    基于 YOLOv8n-cls + 面部关键点
    
    论文: Front. Future Transp., 2025
    """
    
    def __init__(self):
        self.detector = FaceMeshDetector(maxFaces=1)
        # 闭眼阈值：关键点比值 < 18 判定为闭眼
        self.EYE_CLOSED_THRESHOLD = 18
        # 打哈欠阈值：嘴部关键点比值 > 35
        self.YAWN_THRESHOLD = 35
        # 闭眼持续时间阈值（ms）：833ms = 约25帧@30fps
        self.EYE_CLOSED_DURATION_MS = 833
        self.eye_closed_start = None
        self.fatigue_detected = False
        
    def calculate_ear(self, eye_landmarks):
        """
        计算眼睛纵横比（Eye Aspect Ratio的变体）
        
        Args:
            eye_landmarks: 4个眼睛关键点 [(x1,y1), (x2,y2), (x3,y3), (x4,y4)]
            
        Returns:
            ratio: 眼睛开度比值
        """
        # 水平距离
        h_dist = ((eye_landmarks[0][0] - eye_landmarks[3][0])**2 + 
                  (eye_landmarks[0][1] - eye_landmarks[3][1])**2) ** 0.5
        # 垂直距离
        v_dist = ((eye_landmarks[1][1] - eye_landmarks[2][1]))
        
        if v_dist == 0:
            return 999  # 避免除零
        
        ratio = h_dist / v_dist
        return ratio
    
    def process_frame(self, frame):
        """
        处理单帧视频
        
        Args:
            frame: BGR图像
            
        Returns:
            dict: {
                'eye_state': 'open' | 'closed',
                'fatigue': bool,
                'ear': float
            }
        """
        img, faces = self.detector.findFaceMesh(frame, draw=False)
        
        if not faces:
            return {'eye_state': 'no_face', 'fatigue': False, 'ear': 0}
        
        face = faces[0]
        
        # 左眼关键点（MediaPipe索引）
        left_eye = [face[33], face[160], face[158], face[133]]
        right_eye = [face[362], face[385], face[387], face[263]]
        
        left_ear = self.calculate_ear(left_eye)
        right_ear = self.calculate_ear(right_eye)
        avg_ear = (left_ear + right_ear) / 2
        
        eye_closed = avg_ear > self.EYE_CLOSED_THRESHOLD
        
        import time
        current_time = time.time() * 1000  # ms
        
        if eye_closed:
            if self.eye_closed_start is None:
                self.eye_closed_start = current_time
            duration = current_time - self.eye_closed_start
            if duration >= self.EYE_CLOSED_DURATION_MS:
                self.fatigue_detected = True
                return {
                    'eye_state': 'closed',
                    'fatigue': True,
                    'ear': avg_ear,
                    'closed_duration_ms': duration
                }
        else:
            self.eye_closed_start = None
            self.fatigue_detected = False
            
        return {
            'eye_state': 'closed' if eye_closed else 'open',
            'fatigue': self.fatigue_detected,
            'ear': avg_ear
        }


# 实际测试代码
if __name__ == "__main__":
    import numpy as np
    
    detector = RailwayFatigueDetector()
    
    # 模拟铁路驾驶视频处理
    cap = cv2.VideoCapture("railway_driver_test.mp4")
    
    fatigue_events = []
    frame_count = 0
    fps = 30
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        result = detector.process_frame(frame)
        frame_count += 1
        
        if result.get('fatigue'):
            timestamp = frame_count / fps
            fatigue_events.append({
                'time_s': timestamp,
                'ear': result['ear'],
                'closed_ms': result.get('closed_duration_ms', 0)
            })
            print(f"[{timestamp:.1f}s] WARN: 疲劳检测 - 闭眼{result.get('closed_duration_ms', 0):.0f}ms")
    
    cap.release()
    
    print(f"\n=== 检测结果 ===")
    print(f"总帧数: {frame_count}")
    print(f"疲劳事件: {len(fatigue_events)} 次")
    print(f"平均EAR: {np.mean([e['ear'] for e in fatigue_events]):.2f}" if fatigue_events else "无疲劳事件")

from ultralytics import YOLO
import torch

def train_railway_drowsiness_model():
    """
    YOLOv8 铁路驾驶员疲劳分类模型训练
    
    数据集: 6,991张图片 (Awake/Drowsy二分类)
    训练环境: Python 3.8, PyTorch, NVIDIA A100 40GB
    """
    # 加载预训练模型
    model = YOLO('yolov8n-cls.pt')  # nano版本，2.7M参数
    
    # 训练参数（按论文Table 1）
    results = model.train(
        data='railway_drowsiness_dataset/',
        epochs=500,
        batch=1024,
        imgsz=224,
        lr0=0.01,
        optimizer='AdamW',
        # 正则化
        weight_decay=0.01,
        label_smoothing=0.1,
        # 数据增强
        hsv_h=0.015,  # 色调
        hsv_s=0.7,    # 饱和度
        hsv_v=0.4,    # 明度
        degrees=0,    # 旋转（论文用±15°，YOLO内置）
        translate=0.1,
        scale=0.5,
        fliplr=0.5,
        # 早停
        patience=50,
        # 保存
        save=True,
        project='railway_drowsiness',
        name='yolov8n_cls_v1'
    )
    
    # 验证
    metrics = model.val()
    print(f"Accuracy: {metrics.accuracy:.4f}")
    print(f"Precision: {metrics.precision:.4f}")
    print(f"Recall: {metrics.recall:.4f}")
    
    return model

# 导出ONNX用于部署
def export_for_deployment(model_path):
    model = YOLO(model_path)
    # 导出ONNX
    model.export(format='onnx', imgsz=224)
    # 导出INT8量化模型
    model.export(format='engine', imgsz=224, int8=True)
    print("导出完成: ONNX + TensorRT INT8")

import matplotlib.pyplot as plt
import numpy as np

# 按时段的准确率对比
shifts = ['Morning', 'Afternoon', 'Night']
accuracy = [97.13, 95.41, 97.87]
precision = [97.71, 96.61, 97.51]
recall = [97.33, 96.20, 98.85]

fig, ax = plt.subplots(figsize=(10, 6))
x = np.arange(len(shifts))
width = 0.25

ax.bar(x - width, accuracy, width, label='Accuracy', color='#2196F3')
ax.bar(x, precision, width, label='Precision', color='#4CAF50')
ax.bar(x + width, recall, width, label='Recall', color='#FF9800')

ax.set_ylabel('Percentage (%)')
ax.set_title('Railway Driver Drowsiness Detection Performance by Shift')
ax.set_xticks(x)
ax.set_xticklabels(shifts)
ax.legend()
ax.set_ylim(94, 100)

# 标注午班下降原因
ax.annotate('Solar glare\n−2.5% accuracy', 
            xy=(1, 95.41), xytext=(1.3, 94.5),
            arrowprops=dict(arrowstyle='->', color='red'),
            fontsize=9, color='red')
plt.tight_layout()
plt.savefig('railway_drowsiness_by_shift.png', dpi=150)
plt.show()

# IMS疲劳检测阈值验证
THRESHOLDS = {
    'eye_closed_ms': 833,        # 论文验证有效
    'perclos_window_s': 60,      # 60秒滑动窗口
    'perclos_warn_pct': 30,      # 30%闭眼→一级警告
    'perclos_danger_pct': 50,   # 50%闭眼→二级警告
    'yawn_count_threshold': 3,  # 3次/分钟→疲劳
    'blink_rate_threshold': 25  # >25次/分钟→疲劳
}

def validate_fatigue_threshold(eye_openness_series, fps=30):
    """
    按论文标准验证疲劳阈值
    
    Args:
        eye_openness_series: 眼睑开度序列 (0-1)
        fps: 帧率
    
    Returns:
        fatigue_level: 0(清醒) | 1(轻度) | 2(中度) | 3(重度)
    """
    window_frames = 60 * fps  # 60秒窗口
    threshold_ms = 833
    threshold_frames = int(threshold_ms / 1000 * fps)
    
    # 统计闭眼事件
    closed_events = []
    current_closed = 0
    
    for val in eye_openness_series:
        if val < 0.2:  # 闭眼
            current_closed += 1
        else:
            if current_closed >= threshold_frames:
                closed_events.append(current_closed / fps * 1000)
            current_closed = 0
    
    # 计算PERCLOS
    window = eye_openness_series[-window_frames:]
    perclos = np.sum(window < 0.2) / window_frames * 100
    
    # 判定等级
    if perclos >= 50 or len(closed_events) >= 3:
        return 3  # 重度疲劳
    elif perclos >= 30 or len(closed_events) >= 1:
        return 2  # 中度疲劳
    elif perclos >= 15:
        return 1  # 轻度疲劳
    else:
        return 0  # 清醒