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
| """
InCarEmo 多模态特征提取管道
============================
论文 Section 3 & 4
模态编码器:
- 视觉: ViT-Base (RGB) + CNN (IR)
- 音频: Wav2Vec2 微调
- 文本: BERT 中文
"""
import torch
import torch.nn as nn
from transformers import ViTModel, Wav2Vec2Model, BertModel
class InCarEmoFeatureExtractor(nn.Module):
"""
多模态特征提取器
对应论文 Section 4 的特征提取管道
Args:
vit_model: 预训练 ViT 模型名
wav2vec_model: 预训练 Wav2Vec2 模型名
bert_model: 预训练 BERT 模型名
"""
def __init__(
self,
vit_model: str = "google/vit-base-patch16-224",
wav2vec_model: str = "facebook/wav2vec2-base",
bert_model: str = "bert-base-chinese"
):
super().__init__()
self.rgb_encoder = ViTModel.from_pretrained(vit_model)
self.rgb_proj = nn.Linear(768, 256)
self.ir_encoder = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.AdaptiveAvgPool2d((1, 1)),
nn.Flatten(),
nn.Linear(64, 256)
)
self.audio_encoder = Wav2Vec2Model.from_pretrained(wav2vec_model)
self.audio_proj = nn.Linear(768, 256)
self.text_encoder = BertModel.from_pretrained(bert_model)
self.text_proj = nn.Linear(768, 256)
self.modality_weights = nn.Parameter(torch.ones(4))
def forward(
self,
rgb_frames: torch.Tensor = None,
ir_frames: torch.Tensor = None,
audio: torch.Tensor = None,
text_input_ids: torch.Tensor = None,
text_attention_mask: torch.Tensor = None
) -> dict:
"""
提取多模态特征
Args:
rgb_frames: [B, T_rgb, C, H, W] RGB 视频帧
ir_frames: [B, T_ir, 1, H_ir, W_ir] 红外视频帧
audio: [B, T_audio] 原始音频波形
text_input_ids: [B, T_text] 文本 token
text_attention_mask: [B, T_text] 注意力掩码
Returns:
features: dict of modality features
"""
features = {}
if rgb_frames is not None:
B, T, C, H, W = rgb_frames.shape
rgb_flat = rgb_frames.view(B * T, C, H, W)
rgb_out = self.rgb_encoder(rgb_flat).last_hidden_state[:, 0]
rgb_feat = self.rgb_proj(rgb_out).view(B, T, 256).mean(dim=1)
features['rgb'] = rgb_feat * torch.sigmoid(self.modality_weights[0])
if ir_frames is not None:
B, T, _, H_ir, W_ir = ir_frames.shape
ir_flat = ir_frames.view(B * T, 1, H_ir, W_ir)
ir_feat = self.ir_encoder(ir_flat).view(B, T, 256).mean(dim=1)
features['ir'] = ir_feat * torch.sigmoid(self.modality_weights[1])
if audio is not None:
audio_out = self.audio_encoder(audio).last_hidden_state
audio_feat = self.audio_proj(audio_out).mean(dim=1)
features['audio'] = audio_feat * torch.sigmoid(self.modality_weights[2])
if text_input_ids is not None:
text_out = self.text_encoder(
input_ids=text_input_ids,
attention_mask=text_attention_mask
).last_hidden_state[:, 0]
text_feat = self.text_proj(text_out)
features['text'] = text_feat * torch.sigmoid(self.modality_weights[3])
return features
if __name__ == "__main__":
print("InCarEmo 多模态特征提取器")
print("=" * 50)
print("模型组件:")
print(" - RGB: ViT-Base (768→256)")
print(" - IR: 轻量CNN (→256)")
print(" - Audio: Wav2Vec2 (768→256)")
print(" - Text: BERT-Chinese (768→256)")
print()
rgb = torch.randn(2, 10, 3, 224, 224)
ir = torch.randn(2, 10, 1, 224, 224)
audio = torch.randn(2, 16000 * 3)
text_ids = torch.randint(0, 21128, (2, 32))
text_mask = torch.ones(2, 32)
print("预期输出:")
print(" rgb: [2, 256]")
print(" ir: [2, 256]")
print(" audio: [2, 256]")
print(" text: [2, 256]")
|
