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
| import torch
import torch.nn as nn
class TransGazeObject(nn.Module):
"""
TransGaze-Object: 直接预测注视物体类别
输入: 驾驶员面部 + 场景图像
输出: 注视物体类别 + 置信度
物体类别: 道路/后视镜/仪表盘/中控屏/手机/乘客/窗外/其他
"""
GAZE_OBJECTS = [
'road', 'left_mirror', 'right_mirror',
'instrument_cluster', 'center_display',
'phone', 'passenger_seat', 'window', 'other'
]
def __init__(self, n_classes: int = 9, d_model: int = 256):
super().__init__()
self.face_patch = nn.Conv2d(3, d_model, 16, stride=16)
self.face_pos = nn.Parameter(torch.randn(1, 64, d_model))
self.face_encoder = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model, nhead=8, batch_first=True),
num_layers=4
)
self.scene_patch = nn.Conv2d(3, d_model, 16, stride=16)
self.scene_pos = nn.Parameter(torch.randn(1, 256, d_model))
self.scene_encoder = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model, nhead=8, batch_first=True),
num_layers=4
)
self.cross_attn = nn.MultiheadAttention(
d_model, num_heads=8, batch_first=True
)
self.classifier = nn.Sequential(
nn.Linear(d_model, 128), nn.ReLU(),
nn.Linear(128, n_classes)
)
def forward(self, face: torch.Tensor, scene: torch.Tensor) -> dict:
"""
Args:
face: (B, 3, 128, 128)
scene: (B, 3, 256, 256)
"""
face_tokens = self.face_patch(face).flatten(2).permute(0, 2, 1)
face_tokens = face_tokens + self.face_pos[:, :face_tokens.size(1)]
face_tokens = self.face_encoder(face_tokens)
scene_tokens = self.scene_patch(scene).flatten(2).permute(0, 2, 1)
scene_tokens = scene_tokens + self.scene_pos[:, :scene_tokens.size(1)]
scene_tokens = self.scene_encoder(scene_tokens)
gaze_features, _ = self.cross_attn(
face_tokens, scene_tokens, scene_tokens
)
pooled = gaze_features.mean(dim=1)
logits = self.classifier(pooled)
return {
'gaze_object': logits,
'attention': _,
'predicted_object': TransGazeObject.GAZE_OBJECTS
}
if __name__ == "__main__":
model = TransGazeObject(n_classes=9)
face = torch.randn(4, 3, 128, 128)
scene = torch.randn(4, 3, 256, 256)
out = model(face, scene)
print(f"注视物体: {out['gaze_object'].shape}")
print(f"类别: {out['predicted_object']}")
|
