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
| """ 雷达-摄像头前融合方案 在特征提取前进行数据融合 """
import numpy as np import torch import torch.nn as nn import torch.nn.functional as F
class RadarCameraFusion(nn.Module): """ 雷达-摄像头融合网络 架构: - 摄像头分支:CNN特征提取 - 雷达分支:PointNet处理点云 - 融合层:跨模态注意力 - 输出:乘员状态分类 """ def __init__(self, config: dict): """ Args: config: 配置参数 - image_size: 图像尺寸 (H, W) - radar_points: 雷达点云数量 - num_classes: 分类类别数 """ super().__init__() self.image_size = config.get('image_size', (240, 320)) self.radar_points = config.get('radar_points', 1000) self.num_classes = config.get('num_classes', 5) self.camera_encoder = nn.Sequential( nn.Conv2d(3, 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.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 256) ) self.radar_encoder = nn.Sequential( nn.Linear(5, 64), nn.ReLU(), nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 256) ) self.cross_attention = nn.MultiheadAttention( embed_dim=256, num_heads=8, dropout=0.1, batch_first=True ) self.classifier = nn.Sequential( nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, self.num_classes) ) def forward( self, image: torch.Tensor, radar_points: torch.Tensor ) -> torch.Tensor: """ 前向传播 Args: image: 图像 (B, 3, H, W) radar_points: 雷达点云 (B, N, 5) - x, y, z, v, snr Returns: logits: 分类输出 (B, num_classes) """ B = image.shape[0] cam_features = self.camera_encoder(image) radar_features = self.radar_encoder(radar_points) radar_global = radar_features.mean(dim=1) cam_features = cam_features.unsqueeze(1) radar_global = radar_global.unsqueeze(1) fused_features, _ = self.cross_attention( cam_features, radar_global, radar_global ) fused_features = fused_features.squeeze(1) logits = self.classifier(fused_features) return logits
if __name__ == "__main__": config = { 'image_size': (240, 320), 'radar_points': 500, 'num_classes': 5 } model = RadarCameraFusion(config) model.eval() B = 2 image = torch.randn(B, 3, 240, 320) radar_points = torch.randn(B, 500, 5) with torch.no_grad(): logits = model(image, radar_points) print(f"输入图像形状: {image.shape}") print(f"输入雷达点云形状: {radar_points.shape}") print(f"输出logits形状: {logits.shape}") print(f"预测类别: {torch.argmax(logits, dim=-1)}") total_params = sum(p.numel() for p in model.parameters()) print(f"\n模型参数量: {total_params:,} ({total_params/1e6:.2f}M)")
|