-
Notifications
You must be signed in to change notification settings - Fork 0
/
CocoPoseNet.py
262 lines (242 loc) · 15.5 KB
/
CocoPoseNet.py
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
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
import chainer
import chainer.functions as F
import chainer.links as L
from chainer.links import caffe
def copy_vgg_params(model):
print('Copying params of pretrained model...')
layer_names = [
"conv1_1", "conv1_2", "conv2_1", "conv2_2", "conv3_1",
"conv3_2", "conv3_3", "conv3_4", "conv4_1", "conv4_2",
]
pre_model = caffe.CaffeFunction('models/VGG_ILSVRC_19_layers.caffemodel')
for layer_name in layer_names:
exec("model.%s.W.data = pre_model['%s'].W.data" % (layer_name, layer_name))
exec("model.%s.b.data = pre_model['%s'].b.data" % (layer_name, layer_name))
print('Done.')
class CocoPoseNet(chainer.Chain):
insize = 368
def __init__(self):
super(CocoPoseNet, self).__init__(
# cnn to make feature map
conv1_1=L.Convolution2D(in_channels=3, out_channels=64, ksize=3, stride=1, pad=1),
conv1_2=L.Convolution2D(in_channels=64, out_channels=64, ksize=3, stride=1, pad=1),
conv2_1=L.Convolution2D(in_channels=64, out_channels=128, ksize=3, stride=1, pad=1),
conv2_2=L.Convolution2D(in_channels=128, out_channels=128, ksize=3, stride=1, pad=1),
conv3_1=L.Convolution2D(in_channels=128, out_channels=256, ksize=3, stride=1, pad=1),
conv3_2=L.Convolution2D(in_channels=256, out_channels=256, ksize=3, stride=1, pad=1),
conv3_3=L.Convolution2D(in_channels=256, out_channels=256, ksize=3, stride=1, pad=1),
conv3_4=L.Convolution2D(in_channels=256, out_channels=256, ksize=3, stride=1, pad=1),
conv4_1=L.Convolution2D(in_channels=256, out_channels=512, ksize=3, stride=1, pad=1),
conv4_2=L.Convolution2D(in_channels=512, out_channels=512, ksize=3, stride=1, pad=1),
conv4_3_CPM=L.Convolution2D(in_channels=512, out_channels=256, ksize=3, stride=1, pad=1),
conv4_4_CPM=L.Convolution2D(in_channels=256, out_channels=128, ksize=3, stride=1, pad=1),
# stage1
conv5_1_CPM_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=3, stride=1, pad=1),
conv5_2_CPM_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=3, stride=1, pad=1),
conv5_3_CPM_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=3, stride=1, pad=1),
conv5_4_CPM_L1=L.Convolution2D(in_channels=128, out_channels=512, ksize=1, stride=1, pad=0),
conv5_5_CPM_L1=L.Convolution2D(in_channels=512, out_channels=38, ksize=1, stride=1, pad=0),
conv5_1_CPM_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=3, stride=1, pad=1),
conv5_2_CPM_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=3, stride=1, pad=1),
conv5_3_CPM_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=3, stride=1, pad=1),
conv5_4_CPM_L2=L.Convolution2D(in_channels=128, out_channels=512, ksize=1, stride=1, pad=0),
conv5_5_CPM_L2=L.Convolution2D(in_channels=512, out_channels=19, ksize=1, stride=1, pad=0),
# stage2
Mconv1_stage2_L1=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage2_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage2_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage2_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage2_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage2_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage2_L1=L.Convolution2D(in_channels=128, out_channels=38, ksize=1, stride=1, pad=0),
Mconv1_stage2_L2=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage2_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage2_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage2_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage2_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage2_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage2_L2=L.Convolution2D(in_channels=128, out_channels=19, ksize=1, stride=1, pad=0),
# stage3
Mconv1_stage3_L1=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage3_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage3_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage3_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage3_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage3_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage3_L1=L.Convolution2D(in_channels=128, out_channels=38, ksize=1, stride=1, pad=0),
Mconv1_stage3_L2=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage3_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage3_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage3_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage3_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage3_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage3_L2=L.Convolution2D(in_channels=128, out_channels=19, ksize=1, stride=1, pad=0),
# stage4
Mconv1_stage4_L1=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage4_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage4_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage4_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage4_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage4_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage4_L1=L.Convolution2D(in_channels=128, out_channels=38, ksize=1, stride=1, pad=0),
Mconv1_stage4_L2=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage4_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage4_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage4_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage4_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage4_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage4_L2=L.Convolution2D(in_channels=128, out_channels=19, ksize=1, stride=1, pad=0),
# stage5
Mconv1_stage5_L1=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage5_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage5_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage5_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage5_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage5_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage5_L1=L.Convolution2D(in_channels=128, out_channels=38, ksize=1, stride=1, pad=0),
Mconv1_stage5_L2=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage5_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage5_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage5_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage5_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage5_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage5_L2=L.Convolution2D(in_channels=128, out_channels=19, ksize=1, stride=1, pad=0),
# stage6
Mconv1_stage6_L1=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage6_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage6_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage6_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage6_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage6_L1=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage6_L1=L.Convolution2D(in_channels=128, out_channels=38, ksize=1, stride=1, pad=0),
Mconv1_stage6_L2=L.Convolution2D(in_channels=185, out_channels=128, ksize=7, stride=1, pad=3),
Mconv2_stage6_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv3_stage6_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv4_stage6_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv5_stage6_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=7, stride=1, pad=3),
Mconv6_stage6_L2=L.Convolution2D(in_channels=128, out_channels=128, ksize=1, stride=1, pad=0),
Mconv7_stage6_L2=L.Convolution2D(in_channels=128, out_channels=19, ksize=1, stride=1, pad=0),
)
def __call__(self, x):
heatmaps = []
pafs = []
h = F.relu(self.conv1_1(x))
h = F.relu(self.conv1_2(h))
h = F.max_pooling_2d(h, ksize=2, stride=2)
h = F.relu(self.conv2_1(h))
h = F.relu(self.conv2_2(h))
h = F.max_pooling_2d(h, ksize=2, stride=2)
h = F.relu(self.conv3_1(h))
h = F.relu(self.conv3_2(h))
h = F.relu(self.conv3_3(h))
h = F.relu(self.conv3_4(h))
h = F.max_pooling_2d(h, ksize=2, stride=2)
h = F.relu(self.conv4_1(h))
h = F.relu(self.conv4_2(h))
h = F.relu(self.conv4_3_CPM(h))
h = F.relu(self.conv4_4_CPM(h))
feature_map = h
# stage1
h1 = F.relu(self.conv5_1_CPM_L1(feature_map)) # branch1
h1 = F.relu(self.conv5_2_CPM_L1(h1))
h1 = F.relu(self.conv5_3_CPM_L1(h1))
h1 = F.relu(self.conv5_4_CPM_L1(h1))
h1 = self.conv5_5_CPM_L1(h1)
h2 = F.relu(self.conv5_1_CPM_L2(feature_map)) # branch2
h2 = F.relu(self.conv5_2_CPM_L2(h2))
h2 = F.relu(self.conv5_3_CPM_L2(h2))
h2 = F.relu(self.conv5_4_CPM_L2(h2))
h2 = self.conv5_5_CPM_L2(h2)
pafs.append(h1)
heatmaps.append(h2)
# stage2
h = F.concat((h1, h2, feature_map), axis=1) # channel concat
h1 = F.relu(self.Mconv1_stage2_L1(h)) # branch1
h1 = F.relu(self.Mconv2_stage2_L1(h1))
h1 = F.relu(self.Mconv3_stage2_L1(h1))
h1 = F.relu(self.Mconv4_stage2_L1(h1))
h1 = F.relu(self.Mconv5_stage2_L1(h1))
h1 = F.relu(self.Mconv6_stage2_L1(h1))
h1 = self.Mconv7_stage2_L1(h1)
h2 = F.relu(self.Mconv1_stage2_L2(h)) # branch2
h2 = F.relu(self.Mconv2_stage2_L2(h2))
h2 = F.relu(self.Mconv3_stage2_L2(h2))
h2 = F.relu(self.Mconv4_stage2_L2(h2))
h2 = F.relu(self.Mconv5_stage2_L2(h2))
h2 = F.relu(self.Mconv6_stage2_L2(h2))
h2 = self.Mconv7_stage2_L2(h2)
pafs.append(h1)
heatmaps.append(h2)
# stage3
h = F.concat((h1, h2, feature_map), axis=1) # channel concat
h1 = F.relu(self.Mconv1_stage3_L1(h)) # branch1
h1 = F.relu(self.Mconv2_stage3_L1(h1))
h1 = F.relu(self.Mconv3_stage3_L1(h1))
h1 = F.relu(self.Mconv4_stage3_L1(h1))
h1 = F.relu(self.Mconv5_stage3_L1(h1))
h1 = F.relu(self.Mconv6_stage3_L1(h1))
h1 = self.Mconv7_stage3_L1(h1)
h2 = F.relu(self.Mconv1_stage3_L2(h)) # branch2
h2 = F.relu(self.Mconv2_stage3_L2(h2))
h2 = F.relu(self.Mconv3_stage3_L2(h2))
h2 = F.relu(self.Mconv4_stage3_L2(h2))
h2 = F.relu(self.Mconv5_stage3_L2(h2))
h2 = F.relu(self.Mconv6_stage3_L2(h2))
h2 = self.Mconv7_stage3_L2(h2)
pafs.append(h1)
heatmaps.append(h2)
# stage4
h = F.concat((h1, h2, feature_map), axis=1) # channel concat
h1 = F.relu(self.Mconv1_stage4_L1(h)) # branch1
h1 = F.relu(self.Mconv2_stage4_L1(h1))
h1 = F.relu(self.Mconv3_stage4_L1(h1))
h1 = F.relu(self.Mconv4_stage4_L1(h1))
h1 = F.relu(self.Mconv5_stage4_L1(h1))
h1 = F.relu(self.Mconv6_stage4_L1(h1))
h1 = self.Mconv7_stage4_L1(h1)
h2 = F.relu(self.Mconv1_stage4_L2(h)) # branch2
h2 = F.relu(self.Mconv2_stage4_L2(h2))
h2 = F.relu(self.Mconv3_stage4_L2(h2))
h2 = F.relu(self.Mconv4_stage4_L2(h2))
h2 = F.relu(self.Mconv5_stage4_L2(h2))
h2 = F.relu(self.Mconv6_stage4_L2(h2))
h2 = self.Mconv7_stage4_L2(h2)
pafs.append(h1)
heatmaps.append(h2)
# stage5
h = F.concat((h1, h2, feature_map), axis=1) # channel concat
h1 = F.relu(self.Mconv1_stage5_L1(h)) # branch1
h1 = F.relu(self.Mconv2_stage5_L1(h1))
h1 = F.relu(self.Mconv3_stage5_L1(h1))
h1 = F.relu(self.Mconv4_stage5_L1(h1))
h1 = F.relu(self.Mconv5_stage5_L1(h1))
h1 = F.relu(self.Mconv6_stage5_L1(h1))
h1 = self.Mconv7_stage5_L1(h1)
h2 = F.relu(self.Mconv1_stage5_L2(h)) # branch2
h2 = F.relu(self.Mconv2_stage5_L2(h2))
h2 = F.relu(self.Mconv3_stage5_L2(h2))
h2 = F.relu(self.Mconv4_stage5_L2(h2))
h2 = F.relu(self.Mconv5_stage5_L2(h2))
h2 = F.relu(self.Mconv6_stage5_L2(h2))
h2 = self.Mconv7_stage5_L2(h2)
pafs.append(h1)
heatmaps.append(h2)
# stage6
h = F.concat((h1, h2, feature_map), axis=1) # channel concat
h1 = F.relu(self.Mconv1_stage6_L1(h)) # branch1
h1 = F.relu(self.Mconv2_stage6_L1(h1))
h1 = F.relu(self.Mconv3_stage6_L1(h1))
h1 = F.relu(self.Mconv4_stage6_L1(h1))
h1 = F.relu(self.Mconv5_stage6_L1(h1))
h1 = F.relu(self.Mconv6_stage6_L1(h1))
h1 = self.Mconv7_stage6_L1(h1)
h2 = F.relu(self.Mconv1_stage6_L2(h)) # branch2
h2 = F.relu(self.Mconv2_stage6_L2(h2))
h2 = F.relu(self.Mconv3_stage6_L2(h2))
h2 = F.relu(self.Mconv4_stage6_L2(h2))
h2 = F.relu(self.Mconv5_stage6_L2(h2))
h2 = F.relu(self.Mconv6_stage6_L2(h2))
h2 = self.Mconv7_stage6_L2(h2)
pafs.append(h1)
heatmaps.append(h2)
return pafs, heatmaps