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
| """ 视觉乘员分类
基于摄像头图像分类 """
import torch import torch.nn as nn
class VisualOccupantClassifier(nn.Module): """ 视觉乘员分类器 分类:成人/儿童/儿童座椅/空座 """ def __init__(self, num_classes: int = 4): super().__init__() self.backbone = nn.Sequential( nn.Conv2d(3, 32, 3, 2, 1), nn.ReLU(), nn.Conv2d(32, 64, 3, 2, 1), nn.ReLU(), nn.Conv2d(64, 128, 3, 2, 1), nn.ReLU(), nn.AdaptiveAvgPool2d(1) ) self.classifier = nn.Sequential( nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, num_classes) ) self.position_head = nn.Sequential( nn.Linear(128, 32), nn.ReLU(), nn.Linear(32, 2) ) def forward(self, image: torch.Tensor) -> dict: """ Args: image: (B, 3, H, W) Returns: { 'class_logits': (B, num_classes), 'position': (B, 2) } """ features = self.backbone(image).flatten(1) class_logits = self.classifier(features) position = torch.sigmoid(self.position_head(features)) return { 'class_logits': class_logits, 'position': position }
class MultiModalOccupantClassifier: """ 多模态乘员分类器 融合压力传感器 + 视觉 """ def __init__(self): self.pressure_classifier = PressureSensorOccupantClassifier() self.visual_classifier = VisualOccupantClassifier() def classify( self, pressure_map: np.ndarray, image: np.ndarray ) -> dict: """ 多模态分类 Args: pressure_map: 压力分布 image: 图像 Returns: 综合分类结果 """ pressure_result = self.pressure_classifier.classify(pressure_map) with torch.no_grad(): image_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float() / 255.0 visual_result = self.visual_classifier(image_tensor) visual_class = visual_result['class_logits'].argmax(dim=1).item() class_names = ['empty', 'adult', 'child', 'child_seat'] final_class = pressure_result['occupant_type'] if class_names[visual_class] == 'child_seat': final_class = 'child_seat' return { 'final_class': final_class, 'pressure_class': pressure_result['occupant_type'], 'visual_class': class_names[visual_class], 'weight_estimate': pressure_result['weight_estimate'], 'position': pressure_result['position'] }
|