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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
| """
BiFuseNet: 双模态融合的驾驶员损伤检测系统
论文: ECU BMVC25 + ICMI25
核心: RGB + IR 双输入融合，3D面部几何提取
依赖: torch>=2.0, torchvision>=0.15, mediapipe>=0.10
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
class BiFuseNetBackbone(nn.Module):
"""
双模态融合骨干网络
输入: RGB视频帧 + IR红外视频帧
输出: 面部3D几何特征 + 损伤状态估计
"""
def __init__(self, num_classes: int = 3):
super().__init__()
self.rgb_encoder = self._build_encoder(in_channels=3)
self.ir_encoder = self._build_encoder(in_channels=1)
self.fusion = nn.Sequential(
nn.Linear(512 * 2, 256),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(256, 128)
)
self.alcohol_head = nn.Linear(128, 3)
self.fatigue_head = nn.Linear(128, 3)
self.emotion_head = nn.Linear(128, 3)
def _build_encoder(self, in_channels: int) -> nn.Module:
"""轻量3D CNN编码器"""
return nn.Sequential(
nn.Conv3d(in_channels, 32, kernel_size=(3,3,3), padding=1),
nn.BatchNorm3d(32),
nn.ReLU(),
nn.MaxPool3d((1,2,2)),
nn.Conv3d(32, 64, kernel_size=(3,3,3), padding=1),
nn.BatchNorm3d(64),
nn.ReLU(),
nn.MaxPool3d((1,2,2)),
nn.Conv3d(64, 128, kernel_size=(3,3,3), padding=1),
nn.BatchNorm3d(128),
nn.ReLU(),
nn.AdaptiveAvgPool3d((1,1,1)),
nn.Flatten(),
nn.Linear(128, 512)
)
def forward(self, rgb: torch.Tensor, ir: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
前向传播
Args:
rgb: RGB视频片段, shape=(B, 3, T, H, W)
ir: IR视频片段, shape=(B, 1, T, H, W)
Returns:
alcohol_pred: 酒精损伤等级 (B, 3)
fatigue_pred: 疲劳等级 (B, 3)
emotion_pred: 情绪等级 (B, 3)
"""
rgb_feat = self.rgb_encoder(rgb)
ir_feat = self.ir_encoder(ir)
fused = torch.cat([rgb_feat, ir_feat], dim=1)
shared_feat = self.fusion(fused)
alcohol_pred = self.alcohol_head(shared_feat)
fatigue_pred = self.fatigue_head(shared_feat)
emotion_pred = self.emotion_head(shared_feat)
return alcohol_pred, fatigue_pred, emotion_pred
class FacialLandmarkExtractor(nn.Module):
"""
面部关键点提取器
提取68个面部关键点，计算:
- 眨眼频率 (EAR, Eye Aspect Ratio)
- 面部肌肉微动作
- 头部姿态估计
"""
def __init__(self):
super().__init__()
self.landmark_det = nn.Sequential(
nn.Conv2d(3, 64, 7, stride=2, padding=3),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d((1,1)),
nn.Flatten(),
nn.Linear(256, 68 * 2)
)
def compute_ear(self, landmarks: torch.Tensor) -> torch.Tensor:
"""计算眼睛纵横比(EAR)"""
left_eye = landmarks[:, 36:42]
right_eye = landmarks[:, 42:48]
def ear(eye):
vertical_1 = torch.norm(eye[:, 1] - eye[:, 5], dim=1)
vertical_2 = torch.norm(eye[:, 2] - eye[:, 4], dim=1)
horizontal = torch.norm(eye[:, 0] - eye[:, 3], dim=1)
return (vertical_1 + vertical_2) / (2.0 * horizontal + 1e-6)
return (ear(left_eye) + ear(right_eye)) / 2.0
if __name__ == "__main__":
model = BiFuseNetBackbone()
landmark_ext = FacialLandmarkExtractor()
B, T, H, W = 1, 16, 224, 224
rgb = torch.randn(B, 3, T, H, W)
ir = torch.randn(B, 1, T, H, W)
alcohol, fatigue, emotion = model(rgb, ir)
print(f"酒精检测: {F.softmax(alcohol, dim=1)}")
print(f" → 预测: {['清醒', '中度', '严重'][alcohol.argmax()]}")
print(f"疲劳检测: {F.softmax(fatigue, dim=1)}")
print(f" → 预测: {['清醒', '轻度', '严重'][fatigue.argmax()]}")
print(f"情绪检测: {F.softmax(emotion, dim=1)}")
print(f" → 预测: {['平静', '愤怒', '攻击性'][emotion.argmax()]}")
face_img = torch.randn(1, 3, 224, 224)
landmarks = landmark_ext.landmark_det(face_img).view(1, 68, 2)
ear = landmark_ext.compute_ear(landmarks)
print(f"EAR值: {ear.item():.4f} (< 0.2 可能闭眼)")
total_params = sum(p.numel() for p in model.parameters())
print(f"\n模型参数量: {total_params:,}")
|
