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
| """ OOP检测算法(基于3D深度摄像头)
关键技术: 1. 深度图人体分割 2. 3D关键点检测 3. 姿态分类(正常/OOP) """
import numpy as np import torch import torch.nn as nn from typing import Dict, Tuple
class OOPDetector(nn.Module): """ OOP异常姿态检测器 输入: - 深度图:(H, W)单通道深度图 - RGB图:(H, W, 3)RGB图像 输出: - OOP类型:正常/站立/跪姿/侧倾/脚放仪表盘 - 置信度:0-1 - 3D关键点:(17, 3)身体关节点 """ def __init__(self, num_classes=5): super().__init__() self.depth_encoder = nn.Sequential( nn.Conv2d(1, 32, 3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(), ) self.rgb_encoder = nn.Sequential( nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(), ) self.fusion = nn.Sequential( nn.Linear(128 * 2, 256), nn.ReLU(), nn.Dropout(0.3), ) self.classifier = nn.Linear(256, num_classes) self.keypoint_regressor = nn.Linear(256, 17 * 3) def forward(self, depth: torch.Tensor, rgb: torch.Tensor) -> Dict[str, torch.Tensor]: """ 前向传播 Args: depth: (B, 1, H, W) rgb: (B, 3, H, W) Returns: output: { 'oop_type': (B,) OOP类型, 'confidence': (B,) 置信度, 'keypoints_3d': (B, 17, 3) 3D关键点 } """ depth_feat = self.depth_encoder(depth) rgb_feat = self.rgb_encoder(rgb) depth_feat = depth_feat.mean([2, 3]) rgb_feat = rgb_feat.mean([2, 3]) fused_feat = self.fusion(torch.cat([depth_feat, rgb_feat], dim=1)) logits = self.classifier(fused_feat) oop_type = torch.argmax(logits, dim=1) confidence = torch.softmax(logits, dim=1).max(dim=1)[0] keypoints_3d = self.keypoint_regressor(fused_feat).view(-1, 17, 3) return { 'oop_type': oop_type, 'confidence': confidence, 'keypoints_3d': keypoints_3d }
if __name__ == "__main__": """ 模拟OOP检测 """ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = OOPDetector(num_classes=5).to(device) print(f"模型参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M") batch_size = 4 depth = torch.randn(batch_size, 1, 240, 320).to(device) rgb = torch.randn(batch_size, 3, 240, 320).to(device) model.eval() with torch.no_grad(): result = model(depth, rgb) oop_names = ['正常坐姿', '站立姿态', '跪姿', '侧倾', '脚放仪表盘'] print("\n" + "=" * 60) print("OOP检测结果") print("=" * 60) for i in range(batch_size): oop_idx = result['oop_type'][i].item() conf = result['confidence'][i].item() print(f"\n样本{i+1}:") print(f" OOP类型: {oop_names[oop_idx]}") print(f" 置信度: {conf:.2%}") print(f" 3D关键点形状: {result['keypoints_3d'][i].shape}")
|