1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
| import torch import torch.nn as nn from typing import Dict
class DeepFatigueDetector(nn.Module): """ 深度学习疲劳检测器 融合多模态特征: - 眼动特征 - 面部表情 - 头部姿态 - 行为模式 """ def __init__(self, eye_feature_dim: int = 32, face_feature_dim: int = 64, head_feature_dim: int = 16, hidden_dim: int = 128): super().__init__() self.eye_encoder = nn.Sequential( nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, eye_feature_dim), nn.ReLU() ) self.face_encoder = nn.Sequential( nn.Linear(17, 64), nn.ReLU(), nn.Linear(64, face_feature_dim), nn.ReLU() ) self.head_encoder = nn.Sequential( nn.Linear(3, 16), nn.ReLU(), nn.Linear(16, head_feature_dim), nn.ReLU() ) self.temporal_encoder = nn.LSTM( input_size=eye_feature_dim + face_feature_dim + head_feature_dim, hidden_size=hidden_dim, num_layers=2, batch_first=True, bidirectional=True ) self.fatigue_classifier = nn.Sequential( nn.Linear(hidden_dim * 2, 64), nn.ReLU(), nn.Dropout(0.3), nn.Linear(64, 4) ) self.perclos_regressor = nn.Sequential( nn.Linear(hidden_dim * 2, 32), nn.ReLU(), nn.Linear(32, 1), nn.Sigmoid() ) def forward(self, eye_features: torch.Tensor, face_features: torch.Tensor, head_features: torch.Tensor) -> Dict[str, torch.Tensor]: """ 前向传播 Args: eye_features: 眼动特征, shape=(batch, seq_len, 10) face_features: 面部特征, shape=(batch, seq_len, 17) head_features: 头部姿态, shape=(batch, seq_len, 3) Returns: outputs: 检测结果 """ batch_size, seq_len = eye_features.shape[:2] eye_encoded = self.eye_encoder(eye_features) face_encoded = self.face_encoder(face_features) head_encoded = self.head_encoder(head_features) fused_features = torch.cat([eye_encoded, face_encoded, head_encoded], dim=-1) temporal_features, _ = self.temporal_encoder(fused_features) last_features = temporal_features[:, -1, :] fatigue_logits = self.fatigue_classifier(last_features) perclos_pred = self.perclos_regressor(last_features) * 100 return { 'fatigue_logits': fatigue_logits, 'fatigue_probs': torch.softmax(fatigue_logits, dim=-1), 'perclos_prediction': perclos_pred }
|