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
| """
热成像姿态估计模型
基于论文方法：在Raspberry Pi 5上运行
无NPU加速，纯CPU/Edge TPU推理
"""
import torch
import torch.nn as nn
import numpy as np
class ThermalPoseEstimator(nn.Module):
"""
热成像人体姿态估计模型
输入：热成像图 (1, 1, H, W) — 单通道温度分布
输出：关键点热力图 (K, H', W') — K个关键点
关键点定义（座舱简化版，13点）：
0: 头顶, 1: 左肩, 2: 右肩, 3: 左肘, 4: 右肘
5: 左腕, 6: 右腕, 7: 躯干中心, 8: 左髋, 9: 右髋
10: 左膝, 11: 右膝, 12: 面部中心
"""
def __init__(self, n_keypoints=13, input_size=(1, 96, 96)):
super().__init__()
self.n_keypoints = n_keypoints
self.backbone = nn.Sequential(
nn.Conv2d(1, 16, 3, stride=2, padding=1),
nn.BatchNorm2d(16),
nn.ReLU6(inplace=True),
nn.Conv2d(16, 32, 3, stride=2, padding=1),
nn.BatchNorm2d(32),
nn.ReLU6(inplace=True),
nn.Conv2d(32, 32, 3, groups=32, padding=1),
nn.Conv2d(32, 64, 1),
nn.BatchNorm2d(64),
nn.ReLU6(inplace=True),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, groups=64, padding=1),
nn.Conv2d(64, 96, 1),
nn.BatchNorm2d(96),
nn.ReLU6(inplace=True),
nn.MaxPool2d(2),
)
self.head = nn.Sequential(
nn.ConvTranspose2d(96, 64, 4, stride=2, padding=1),
nn.ReLU6(inplace=True),
nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
nn.ReLU6(inplace=True),
nn.ConvTranspose2d(32, n_keypoints, 4, stride=2, padding=1),
)
def forward(self, x):
features = self.backbone(x)
heatmaps = self.head(features)
return heatmaps
def predict_keypoints(self, thermal_image):
"""
从热成像图预测关键点
Args:
thermal_image: shape=(1, 1, 96, 96), 归一化温度图
Returns:
keypoints: shape=(13, 2), 各关键点(x,y)坐标
confidences: shape=(13,), 各关键点置信度
"""
with torch.no_grad():
heatmaps = self.forward(thermal_image)
keypoints = np.zeros((self.n_keypoints, 2))
confidences = np.zeros(self.n_keypoints)
for k in range(self.n_keypoints):
hm = heatmaps[0, k].cpu().numpy()
idx = np.unravel_index(np.argmax(hm), hm.shape)
confidences[k] = hm[idx]
scale_x = 96 / hm.shape[1]
scale_y = 96 / hm.shape[0]
keypoints[k] = [idx[1] * scale_x, idx[0] * scale_y]
return keypoints, confidences
class QuantizedThermalPose:
"""量化后的热成像姿态估计模型"""
def __init__(self):
self.model = ThermalPoseEstimator()
self.model_size_fp32 = sum(p.numel() * 4 for p in self.model.parameters()) / 1e6
def quantize_int8(self):
"""INT8量化"""
quantized = torch.quantization.quantize_dynamic(
self.model,
{nn.Conv2d, nn.Linear},
dtype=torch.qint8
)
self.model_size_int8 = sum(
p.element_size() * p.numel()
for p in quantized.parameters()
) / 1e6
return quantized
def benchmark_rpi5(self):
"""
Raspberry Pi 5 性能基准
预期结果（论文参考）：
- FP32: ~150ms/frame, 6.7fps
- INT8: ~45ms/frame, 22fps
- INT8 + Edge TPU: ~15ms/frame, 67fps
"""
results = {
'fp32_cpu': {'latency_ms': 150, 'fps': 6.7, 'size_mb': 0.5},
'int8_cpu': {'latency_ms': 45, 'fps': 22, 'size_mb': 0.15},
'int8_edge_tpu': {'latency_ms': 15, 'fps': 67, 'size_mb': 0.15},
}
return results
|
