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
| """
YOLO26 用于安全带检测的边缘部署方案
架构：
1. YOLO26 目标检测（人体 + 安全带）
2. 分类头判断安全带状态
3. 量化到 INT8 用于 NPU
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class YOLO26BeltDetector(nn.Module):
"""
YOLO26 安全带检测器
基于论文描述的 Edge AI 方案
适用于座舱内摄像头（DMS/OMS）
两阶段方案：
1. YOLO26 检测人体 + 安全带区域
2. 分类头判断安全带佩戴状态
"""
def __init__(self, num_belt_classes=5,
backbone_channels=32):
super().__init__()
self.backbone = self._build_backbone(backbone_channels)
self.detection_head = nn.ModuleList([
self._build_head(backbone_channels*8, 2),
self._build_head(backbone_channels*4, 2),
self._build_head(backbone_channels*2, 2),
])
self.classifier = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(backbone_channels*8, 64),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(64, num_belt_classes)
)
def _build_backbone(self, base_ch):
"""轻量骨干网络"""
layers = []
in_ch = 3
for out_ch in [base_ch, base_ch*2, base_ch*4,
base_ch*8, base_ch*16]:
layers.extend([
nn.Conv2d(in_ch, out_ch, 3, 2, 1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU6(inplace=True),
])
in_ch = out_ch
return nn.Sequential(*layers)
def _build_head(self, in_ch, num_classes):
"""检测头"""
return nn.Sequential(
nn.Conv2d(in_ch, in_ch//2, 1, bias=False),
nn.BatchNorm2d(in_ch//2),
nn.ReLU6(inplace=True),
nn.Conv2d(in_ch//2, num_classes + 4, 1),
)
def forward(self, x):
"""前向传播"""
features = []
for layer in self.backbone:
x = layer(x)
if isinstance(layer, nn.ReLU6):
features.append(x)
detections = []
for i, head in enumerate(self.detection_head):
detections.append(head(features[-(i+1)]))
belt_status = self.classifier(features[-1])
return detections, belt_status
def quantize_model(model):
"""
YOLO26 INT8 量化
适用于 NPU 部署：
- Qualcomm Hexagon NPU
- ARM Ethos NPU
- Intel OpenVINO NPU
"""
model_cpu = model.cpu()
quantized = torch.quantization.quantize_dynamic(
model_cpu,
{nn.Conv2d, nn.Linear},
dtype=torch.quint8
)
original_size = sum(
p.nelement() * p.element_size()
for p in model_cpu.parameters()
)
quant_size = sum(
p.nelement() * p.element_size()
for p in quantized.parameters()
)
return quantized, original_size, quant_size
if __name__ == "__main__":
model = YOLO26BeltDetector(
num_belt_classes=5,
backbone_channels=16
)
x = torch.randn(1, 3, 640, 640)
detections, belt_status = model(x)
print(f"=== YOLO26 安全带检测器 ===")
print(f" 输入: {x.shape}")
print(f" 检测输出: {len(detections)} 个尺度")
for i, det in enumerate(detections):
print(f" Scale {i}: {det.shape}")
print(f" 分类输出: {belt_status.shape}")
print(f" 参数量: {sum(p.numel() for p in model.parameters()):,}")
quant_model, orig_size, quant_size = quantize_model(model)
print(f"\n=== 量化结果 ===")
print(f" 原始大小: {orig_size/1024:.1f} KB")
print(f" 量化大小: {quant_size/1024:.1f} KB")
print(f" 压缩比: {orig_size/quant_size:.1f}x")
params = sum(p.numel() for p in model.parameters())
print(f"\n=== 边缘部署估算 ===")
print(f" 参数量: {params:,}")
print(f" FP32 大小: {params*4/1024:.1f} KB")
print(f" INT8 大小: {params/1024:.1f} KB")
print(f" 预估推理: 5-15 ms (QCS8255 NPU)")
print(f" 预估功耗: <0.5 W")
|
