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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
| import torch
import torch.nn as nn
import torch.nn.functional as F
class RadarDensifier(nn.Module):
"""
Radar Densifier: 解决雷达点云稀疏性
论文核心: 将非空雷达 "pillar" 的特征
用自适应高斯函数重分配到邻域空 pillar
高斯宽度基于 RCS 自适应:
- 大 RCS (大物体) → 宽高斯 (远距离传播)
- 小 RCS (小物体) → 窄高斯 (精细传播)
"""
def __init__(self, in_channels=64, out_channels=64,
grid_size=(200, 200), max_range=50.0):
super().__init__()
self.grid_size = grid_size
self.max_range = max_range
self.feature_encoder = nn.Sequential(
nn.Linear(in_channels + 5, 64),
nn.ReLU(),
nn.Linear(64, out_channels)
)
self.rcs_to_sigma = nn.Sequential(
nn.Linear(1, 16),
nn.ReLU(),
nn.Linear(16, 1),
nn.Softplus()
)
def forward(self, radar_points, pillar_grid=None):
"""
Args:
radar_points: (N, 7+) [x, y, z, rcs, doppler, intensity, ...]
pillar_grid: (G, G, C) 预初始化的 pillar 网格
Returns:
densified_grid: (G, G, C) 致密化后的网格
"""
N = radar_points.shape[0]
G = self.grid_size[0]
x = radar_points[:, 0]
y = radar_points[:, 1]
rcs = radar_points[:, 3]
gx = ((x + self.max_range) / (2 * self.max_range) * G).long().clamp(0, G-1)
gy = ((y + self.max_range) / (2 * self.max_range) * G).long().clamp(0, G-1)
grid = torch.zeros(G, G, 64)
for i in range(N):
if 0 <= gx[i] < G and 0 <= gy[i] < G:
feat = self.feature_encoder(
radar_points[i:i+1]
).squeeze(0)
grid[gx[i], gy[i]] += feat
densified_grid = grid.clone()
sigma_base = self.rcs_to_sigma(rcs.unsqueeze(-1)).squeeze(-1)
non_empty_mask = grid.sum(dim=-1) > 0
non_empty_idx = torch.where(non_empty_mask)
for (px, py) in zip(*non_empty_idx):
sigma = sigma_base[px * G + py] if px * G + py < N else 1.0
sigma = max(sigma.item(), 0.5)
radius = int(3 * sigma)
for dx in range(-radius, radius + 1):
for dy in range(-radius, radius + 1):
nx, ny = px + dx, py + dy
if 0 <= nx < G and 0 <= ny < G:
if dx == 0 and dy == 0:
continue
dist_sq = dx * dx + dy * dy
weight = torch.exp(
torch.tensor(-dist_sq / (2 * sigma ** 2))
)
densified_grid[nx, ny] += weight * grid[px, py] * 0.3
return densified_grid
class RadarAmplifier(nn.Module):
"""
Radar Amplifier: 特征通道选择性增强
估计每个特征通道的重要性概率
增强有用通道，抑制噪声通道
"""
def __init__(self, channels=64, reduction=8):
super().__init__()
self.channel_gate = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels, channels // reduction),
nn.ReLU(),
nn.Linear(channels // reduction, channels),
nn.Sigmoid()
)
self.scale = nn.Parameter(torch.ones(1, channels, 1, 1))
def forward(self, x):
"""
Args:
x: (B, C, H, W) 雷达特征图
Returns:
amplified: (B, C, H, W) 增强后的特征
"""
importance = self.channel_gate(x)
importance = importance.unsqueeze(-1).unsqueeze(-1)
amplified = x * importance * self.scale
return amplified
class REOccFusion(nn.Module):
"""
REOcc: 摄像头-雷达融合 3D 占用预测
架构:
1. 雷达 → Densifier → Amplifier → 增强雷达特征
2. 摄像头 → 2D backbone → BEV 转换 → 摄像头 BEV 特征
3. 可变形交叉注意力融合
4. 3D 占用预测头
"""
def __init__(self, radar_channels=64, cam_channels=256,
fusion_dim=256, num_classes=17):
super().__init__()
self.radar_densifier = RadarDensifier(radar_channels)
self.radar_amplifier = RadarAmplifier(radar_channels)
self.radar_proj = nn.Conv2d(radar_channels, fusion_dim, 1)
self.cam_backbone = nn.Sequential(
nn.Conv2d(3, 64, 7, stride=2, padding=3),
nn.BatchNorm2d(64), nn.ReLU(),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.BatchNorm2d(128), nn.ReLU(),
nn.Conv2d(128, cam_channels, 3, stride=2, padding=1),
nn.BatchNorm2d(cam_channels), nn.ReLU(),
)
self.cam_proj = nn.Conv2d(cam_channels, fusion_dim, 1)
self.fusion = DeformableCrossAttention(
dim=fusion_dim, num_heads=8
)
self.occ_head = nn.Sequential(
nn.Conv3d(fusion_dim, 128, 3, padding=1),
nn.BatchNorm3d(128), nn.ReLU(),
nn.Conv3d(128, num_classes, 1)
)
def forward(self, radar_points, cam_images):
"""
Args:
radar_points: (N, 7) 雷达点
cam_images: (B, 3, H, W) 摄像头图像
Returns:
occ_pred: (B, num_classes, D, H, W) 3D 占用预测
"""
radar_grid = self.radar_densifier(radar_points)
radar_grid = radar_grid.permute(2, 0, 1).unsqueeze(0)
radar_feat = self.radar_amplifier(radar_grid)
radar_feat = self.radar_proj(radar_feat)
cam_feat = self.cam_backbone(cam_images)
cam_feat = self.cam_proj(cam_feat)
fused = self.fusion(cam_feat, radar_feat)
occ = self.occ_head(fused.unsqueeze(-1))
return occ
class DeformableCrossAttention(nn.Module):
"""可变形交叉注意力"""
def __init__(self, dim=256, num_heads=8):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.q_proj = nn.Linear(dim, dim)
self.k_proj = nn.Linear(dim, dim)
self.v_proj = nn.Linear(dim, dim)
self.offset_gen = nn.Linear(dim, num_heads * 2)
self.out_proj = nn.Linear(dim, dim)
self.scale = self.head_dim ** -0.5
def forward(self, query_feat, key_value_feat):
B, C, H, W = query_feat.shape
q = self.q_proj(
query_feat.flatten(2).transpose(1, 2)
).transpose(1, 2).reshape(B, self.num_heads, -1, H*W)
k = self.k_proj(
key_value_feat.flatten(2).transpose(1, 2)
).transpose(1, 2).reshape(B, self.num_heads, -1, H*W)
v = self.v_proj(
key_value_feat.flatten(2).transpose(1, 2)
).transpose(1, 2).reshape(B, self.num_heads, -1, H*W)
attn = torch.matmul(q.transpose(-1, -2), k) * self.scale
attn = torch.softmax(attn, dim=-1)
out = torch.matmul(attn, v.transpose(-1, -2))
out = out.reshape(B, C, H, W)
out = self.out_proj(
out.flatten(2).transpose(1, 2)
).transpose(1, 2).reshape(B, C, H, W)
return query_feat + out
if __name__ == "__main__":
amp = RadarAmplifier(64)
x = torch.randn(2, 64, 200, 200)
out = amp(x)
print(f"RadarAmplifier: {x.shape} → {out.shape}")
print(f"Params: {sum(p.numel() for p in amp.parameters()):,}")
|
