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
| import numpy as np import torch import torch.nn as nn from typing import Dict, Tuple from dataclasses import dataclass from enum import Enum
class PostureType(Enum): """乘坐姿势枚举""" NORMAL = 0 FORWARD_LEAN = 1 RECLINE = 2 LEFT_TILT = 3 RIGHT_TILT = 4 CROSSED_LEGS = 5 ONE_LEG_CROSSED = 6 SIDE_SITTING = 7 SLUMPED = 8
@dataclass class PressureDistribution: """压力分布数据""" pressure_map: np.ndarray timestamp: float seat_type: str
class PressureMapCNN(nn.Module): """压力分布CNN分类器 基于MDPI Applied Sciences论文实现 """ def __init__(self, num_classes: int = 9): super().__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.bn1 = nn.BatchNorm2d(32) self.bn2 = nn.BatchNorm2d(64) self.bn3 = nn.BatchNorm2d(128) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(128 * 2 * 2, 256) self.fc2 = nn.Linear(256, 128) self.fc3 = nn.Linear(128, num_classes) self.dropout = nn.Dropout(0.5) def forward(self, x): """ Args: x: (B, 1, 16, 16) 压力分布图 Returns: logits: (B, 9) 姿势分类 """ x = self.pool(nn.functional.relu(self.bn1(self.conv1(x)))) x = self.pool(nn.functional.relu(self.bn2(self.conv2(x)))) x = self.pool(nn.functional.relu(self.bn3(self.conv3(x)))) x = x.view(x.size(0), -1) x = nn.functional.relu(self.fc1(x)) x = self.dropout(x) x = nn.functional.relu(self.fc2(x)) x = self.dropout(x) x = self.fc3(x) return x
class PostureRecognizer: """乘坐姿势识别器""" def __init__(self, model_path: str = None): self.model = PressureMapCNN() if model_path: self.model.load_state_dict(torch.load(model_path)) self.model.eval() self.oop_risk = { PostureType.NORMAL: 'normal', PostureType.FORWARD_LEAN: 'high', PostureType.RECLINE: 'medium', PostureType.LEFT_TILT: 'medium', PostureType.RIGHT_TILT: 'medium', PostureType.CROSSED_LEGS: 'medium', PostureType.ONE_LEG_CROSSED: 'medium', PostureType.SIDE_SITTING: 'high', PostureType.SLUMPED: 'high' } def recognize(self, pressure_data: PressureDistribution) -> Dict: """ 识别乘坐姿势 Args: pressure_data: 压力分布数据 Returns: result: 识别结果 """ pressure_map = self._preprocess(pressure_data.pressure_map) with torch.no_grad(): input_tensor = torch.FloatTensor(pressure_map).unsqueeze(0).unsqueeze(0) logits = self.model(input_tensor) probs = torch.softmax(logits, dim=1) pred_class = torch.argmax(probs, dim=1).item() posture = PostureType(pred_class) return { 'posture': posture.name, 'confidence': probs[0, pred_class].item(), 'oop_risk': self.oop_risk[posture], 'all_probs': {PostureType(i).name: probs[0, i].item() for i in range(9)} } def _preprocess(self, pressure_map: np.ndarray) -> np.ndarray: """预处理压力图""" pressure_map = (pressure_map - pressure_map.min()) / (pressure_map.max() - pressure_map.min() + 1e-6) if pressure_map.shape != (16, 16): from scipy.ndimage import zoom zoom_factor = (16 / pressure_map.shape[0], 16 / pressure_map.shape[1]) pressure_map = zoom(pressure_map, zoom_factor) return pressure_map
if __name__ == "__main__": recognizer = PostureRecognizer() pressure_map = np.zeros((16, 16)) pressure_map[8:16, :] = np.random.rand(8, 16) * 0.8 + 0.2 pressure_map[0:8, :] = np.random.rand(8, 16) * 0.3 pressure_data = PressureDistribution( pressure_map=pressure_map, timestamp=time.time(), seat_type='soft' ) result = recognizer.recognize(pressure_data) print(f"姿势判定: {result['posture']}") print(f"置信度: {result['confidence']:.2f}") print(f"OOP风险: {result['oop_risk']}") print(f"\n各姿势概率:") for posture, prob in result['all_probs'].items(): print(f" {posture}: {prob:.3f}")
|