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
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
| import torch
import torch.nn as nn
import torch.nn.functional as F
class PolishNet(nn.Module):
"""
PolishNet: 压力图到人体图像的域转换器
论文核心方法实现
架构: 编码-解码"沙漏"全卷积网络
功能: 将模糊的压力传感器热图转换为标准姿态估计器
可理解的人体形状图像
损失函数:
- Pixel Loss: 像素级重建保真
- Heatmap Loss: 关节位置准确性
- PAF Loss: 肢体方向和结构
"""
def __init__(self, in_channels=1, out_channels=3,
base_channels=64, num_res_blocks=3):
super().__init__()
self.encoder = self._make_encoder(in_channels, base_channels)
self.bottleneck = self._make_bottleneck(base_channels * 8)
self.decoder = self._make_decoder(base_channels * 8, out_channels)
self.final_conv = nn.Conv2d(base_channels, out_channels,
kernel_size=7, padding=3)
def _make_encoder(self, in_ch, base_ch):
"""编码器：4层下采样"""
layers = []
channels = [in_ch, base_ch, base_ch*2, base_ch*4, base_ch*8]
for i in range(len(channels)-1):
layers.extend([
nn.Conv2d(channels[i], channels[i+1],
kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(channels[i+1]),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(channels[i+1], channels[i+1],
kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(channels[i+1]),
nn.LeakyReLU(0.2, inplace=True)
])
return nn.Sequential(*layers)
def _make_bottleneck(self, channels):
"""瓶颈层：残差块"""
layers = []
for _ in range(3):
layers.append(ResidualBlock(channels))
return nn.Sequential(*layers)
def _make_decoder(self, in_ch, out_ch):
"""解码器：4层上采样"""
layers = []
channels = [in_ch, in_ch//2, in_ch//4, in_ch//8, in_ch//16]
for i in range(len(channels)-1):
layers.extend([
nn.ConvTranspose2d(channels[i], channels[i+1],
kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(channels[i+1]),
nn.ReLU(inplace=True),
nn.Conv2d(channels[i+1], channels[i+1],
kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(channels[i+1]),
nn.ReLU(inplace=True)
])
return nn.Sequential(*layers)
def forward(self, pressure_map):
"""
Args:
pressure_map: (B, 1, H, W) 压力传感器矩阵
Returns:
polished_image: (B, 3, H, W) 人体形状图像
"""
feat = self.encoder(pressure_map)
feat = self.bottleneck(feat)
img = self.decoder(feat)
img = torch.sigmoid(self.final_conv(img))
return img
class ResidualBlock(nn.Module):
"""残差块"""
def __init__(self, channels):
super().__init__()
self.conv1 = nn.Conv2d(channels, channels, 3, 1, 1)
self.bn1 = nn.BatchNorm2d(channels)
self.conv2 = nn.Conv2d(channels, channels, 3, 1, 1)
self.bn2 = nn.BatchNorm2d(channels)
def forward(self, x):
residual = x
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
return F.relu(out + residual)
class PolishNetLoss(nn.Module):
"""
PolishNet多任务损失函数
论文 Section 3.2 实现
"""
def __init__(self, lambda_pixel=1.0, lambda_heatmap=1.0, lambda_paf=0.5):
super().__init__()
self.lambda_pixel = lambda_pixel
self.lambda_heatmap = lambda_heatmap
self.lambda_paf = lambda_paf
self.l1_loss = nn.L1Loss()
self.mse_loss = nn.MSELoss()
def forward(self, pred_img, target_img,
pred_heatmap=None, target_heatmap=None,
pred_paf=None, target_paf=None):
"""
综合损失计算
"""
pixel_loss = self.l1_loss(pred_img, target_img)
heatmap_loss = 0
if pred_heatmap is not None and target_heatmap is not None:
heatmap_loss = self.mse_loss(pred_heatmap, target_heatmap)
paf_loss = 0
if pred_paf is not None and target_paf is not None:
paf_loss = self.mse_loss(pred_paf, target_paf)
total_loss = (self.lambda_pixel * pixel_loss +
self.lambda_heatmap * heatmap_loss +
self.lambda_paf * paf_loss)
return total_loss, {
'pixel': pixel_loss.item(),
'heatmap': heatmap_loss if isinstance(heatmap_loss, float) else heatmap_loss.item(),
'paf': paf_loss if isinstance(paf_loss, float) else paf_loss.item()
}
class PressureToPosePipeline(nn.Module):
"""
压力传感器 → 姿态估计完整管道
PolishNet + 冻结的OpenPose
"""
def __init__(self):
super().__init__()
self.polishnet = PolishNet(in_channels=1, out_channels=3)
self.pose_estimator = self._load_frozen_openpose()
def _load_frozen_openpose(self):
"""加载冻结的标准姿态估计器"""
estimator = nn.Sequential(
nn.Conv2d(3, 64, 3, 1, 1),
nn.ReLU(),
nn.Conv2d(64, 17, 1)
)
for p in estimator.parameters():
p.requires_grad = False
return estimator
def forward(self, pressure_map):
"""
Args:
pressure_map: (B, 1, H, W)
Returns:
heatmaps: (B, 17, H, W) 关节热图
polished_img: (B, 3, H, W) 中间图像
"""
polished_img = self.polishnet(pressure_map)
heatmaps = self.pose_estimator(polished_img)
return heatmaps, polished_img
if __name__ == "__main__":
model = PressureToPosePipeline()
pressure = torch.randn(4, 1, 64, 64)
heatmaps, polished = model(pressure)
print(f"Pressure input: {pressure.shape}")
print(f"Polished image: {polished.shape}")
print(f"Joint heatmaps: {heatmaps.shape}")
print(f"PolishNet params: {sum(p.numel() for p in model.polishnet.parameters()):,}")
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Trainable: {trainable:,} / Total: {total:,}")
|
