姿态估计openpose_pytorch_code浅析(待补充)
接上文，经过了openpose的原理简单的解析，这一节我们主要进行code的解析。
CODE解析
我们主要参考的代码是https://github.com/tensorboy/pytorch_Realtime_Multi-Person_Pose_Estimation，代码写的很好，我们主要看的是demo/picture_demo.py
首先我们看下效果，
作图表示输入的图片，酷酷的四字弟弟，右图是出来的关键点，我们根据demo中的单张图片的前向过程讲解下模型的inference阶段，在简单的过一下训练的过程。
而在inference的阶段中，我们主要看这几个关键的函数，我们把这几个函数扒出来单独介绍下。其中相对重要的几个函数，我们主要进行了标红，其中权重文件，在这个git里面有相关的下载连接。
咱们主要看这这几部分
1. model = get_model('vgg19')
表示在上面所述的网络的示意图中，那个F使用的是vgg19，提取到的feature maps。此函数在lib/network/rtpose_vgg.py之中。
1 """CPM Pytorch Implementation""" 2 3 from collections import OrderedDict 4 5 import torch 6 import torch.nn as nn 7 import torch.nn.functional as F 8 import torch.utils.data as data 9 import torch.utils.model_zoo as model_zoo 10 from torch.autograd import Variable 11 from torch.nn import init 12 13 def make_stages(cfg_dict): 14 """Builds CPM stages from a dictionary 15 Args: 16 cfg_dict: a dictionary 17 """ 18 layers = [] 19 for i in range(len(cfg_dict) - 1): 20 one_ = cfg_dict[i] 21 for k, v in one_.items(): 22 if 'pool' in k: 23 layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], 24 padding=v[2])] 25 else: 26 conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], 27 kernel_size=v[2], stride=v[3], 28 padding=v[4]) 29 layers += [conv2d, nn.ReLU(inplace=True)] 30 one_ = list(cfg_dict[-1].keys()) 31 k = one_[0] 32 v = cfg_dict[-1][k] 33 conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], 34 kernel_size=v[2], stride=v[3], padding=v[4]) 35 layers += [conv2d] 36 return nn.Sequential(*layers) 37 38 39 def make_vgg19_block(block): 40 """Builds a vgg19 block from a dictionary 41 Args: 42 block: a dictionary 43 """ 44 layers = [] 45 for i in range(len(block)): 46 one_ = block[i] 47 for k, v in one_.items(): 48 if 'pool' in k: 49 layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], 50 padding=v[2])] 51 else: 52 conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], 53 kernel_size=v[2], stride=v[3], 54 padding=v[4]) 55 layers += [conv2d, nn.ReLU(inplace=True)] 56 return nn.Sequential(*layers) 57 58 59 60 def get_model(trunk='vgg19'): 61 """Creates the whole CPM model 62 Args: 63 trunk: string, 'vgg19' or 'mobilenet' 64 Returns: Module, the defined model 65 """ 66 blocks = {} 67 # block0 is the preprocessing stage 68 if trunk == 'vgg19': 69 block0 = [{'conv1_1': [3, 64, 3, 1, 1]}, 70 {'conv1_2': [64, 64, 3, 1, 1]}, 71 {'pool1_stage1': [2, 2, 0]}, 72 {'conv2_1': [64, 128, 3, 1, 1]}, 73 {'conv2_2': [128, 128, 3, 1, 1]}, 74 {'pool2_stage1': [2, 2, 0]}, 75 {'conv3_1': [128, 256, 3, 1, 1]}, 76 {'conv3_2': [256, 256, 3, 1, 1]}, 77 {'conv3_3': [256, 256, 3, 1, 1]}, 78 {'conv3_4': [256, 256, 3, 1, 1]}, 79 {'pool3_stage1': [2, 2, 0]}, 80 {'conv4_1': [256, 512, 3, 1, 1]}, 81 {'conv4_2': [512, 512, 3, 1, 1]}, 82 {'conv4_3_CPM': [512, 256, 3, 1, 1]}, 83 {'conv4_4_CPM': [256, 128, 3, 1, 1]}] 84 85 elif trunk == 'mobilenet': 86 block0 = [{'conv_bn': [3, 32, 2]}, # out: 3, 32, 184, 184 87 {'conv_dw1': [32, 64, 1]}, # out: 32, 64, 184, 184 88 {'conv_dw2': [64, 128, 2]}, # out: 64, 128, 92, 92 89 {'conv_dw3': [128, 128, 1]}, # out: 128, 256, 92, 92 90 {'conv_dw4': [128, 256, 2]}, # out: 256, 256, 46, 46 91 {'conv4_3_CPM': [256, 256, 1, 3, 1]}, 92 {'conv4_4_CPM': [256, 128, 1, 3, 1]}] 93 94 # Stage 1 95 blocks['block1_1'] = [{'conv5_1_CPM_L1': [128, 128, 3, 1, 1]}, 96 {'conv5_2_CPM_L1': [128, 128, 3, 1, 1]}, 97 {'conv5_3_CPM_L1': [128, 128, 3, 1, 1]}, 98 {'conv5_4_CPM_L1': [128, 512, 1, 1, 0]}, 99 {'conv5_5_CPM_L1': [512, 38, 1, 1, 0]}] 100 101 blocks['block1_2'] = [{'conv5_1_CPM_L2': [128, 128, 3, 1, 1]}, 102 {'conv5_2_CPM_L2': [128, 128, 3, 1, 1]}, 103 {'conv5_3_CPM_L2': [128, 128, 3, 1, 1]}, 104 {'conv5_4_CPM_L2': [128, 512, 1, 1, 0]}, 105 {'conv5_5_CPM_L2': [512, 19, 1, 1, 0]}] 106 107 # Stages 2 - 6 108 for i in range(2, 7): 109 blocks['block%d_1' % i] = [ 110 {'Mconv1_stage%d_L1' % i: [185, 128, 7, 1, 3]}, 111 {'Mconv2_stage%d_L1' % i: [128, 128, 7, 1, 3]}, 112 {'Mconv3_stage%d_L1' % i: [128, 128, 7, 1, 3]}, 113 {'Mconv4_stage%d_L1' % i: [128, 128, 7, 1, 3]}, 114 {'Mconv5_stage%d_L1' % i: [128, 128, 7, 1, 3]}, 115 {'Mconv6_stage%d_L1' % i: [128, 128, 1, 1, 0]}, 116 {'Mconv7_stage%d_L1' % i: [128, 38, 1, 1, 0]} 117 ] 118 119 blocks['block%d_2' % i] = [ 120 {'Mconv1_stage%d_L2' % i: [185, 128, 7, 1, 3]}, 121 {'Mconv2_stage%d_L2' % i: [128, 128, 7, 1, 3]}, 122 {'Mconv3_stage%d_L2' % i: [128, 128, 7, 1, 3]}, 123 {'Mconv4_stage%d_L2' % i: [128, 128, 7, 1, 3]}, 124 {'Mconv5_stage%d_L2' % i: [128, 128, 7, 1, 3]}, 125 {'Mconv6_stage%d_L2' % i: [128, 128, 1, 1, 0]}, 126 {'Mconv7_stage%d_L2' % i: [128, 19, 1, 1, 0]} 127 ] 128 129 models = {} 130 131 if trunk == 'vgg19': 132 print("Bulding VGG19") 133 models['block0'] = make_vgg19_block(block0) 134 135 for k, v in blocks.items(): 136 models[k] = make_stages(list(v)) 137 138 class rtpose_model(nn.Module): 139 def __init__(self, model_dict): 140 super(rtpose_model, self).__init__() 141 self.model0 = model_dict['block0'] 142 self.model1_1 = model_dict['block1_1'] 143 self.model2_1 = model_dict['block2_1'] 144 self.model3_1 = model_dict['block3_1'] 145 self.model4_1 = model_dict['block4_1'] 146 self.model5_1 = model_dict['block5_1'] 147 self.model6_1 = model_dict['block6_1'] 148 149 self.model1_2 = model_dict['block1_2'] 150 self.model2_2 = model_dict['block2_2'] 151 self.model3_2 = model_dict['block3_2'] 152 self.model4_2 = model_dict['block4_2'] 153 self.model5_2 = model_dict['block5_2'] 154 self.model6_2 = model_dict['block6_2'] 155 156 self._initialize_weights_norm() 157 158 def forward(self, x): 159 160 saved_for_loss = [] 161 out1 = self.model0(x) 162 163 out1_1 = self.model1_1(out1) 164 out1_2 = self.model1_2(out1) 165 out2 = torch.cat([out1_1, out1_2, out1], 1) 166 saved_for_loss.append(out1_1) 167 saved_for_loss.append(out1_2) 168 169 out2_1 = self.model2_1(out2) 170 out2_2 = self.model2_2(out2) 171 out3 = torch.cat([out2_1, out2_2, out1], 1) 172 saved_for_loss.append(out2_1) 173 saved_for_loss.append(out2_2) 174 175 out3_1 = self.model3_1(out3) 176 out3_2 = self.model3_2(out3) 177 out4 = torch.cat([out3_1, out3_2, out1], 1) 178 saved_for_loss.append(out3_1) 179 saved_for_loss.append(out3_2) 180 181 out4_1 = self.model4_1(out4) 182 out4_2 = self.model4_2(out4) 183 out5 = torch.cat([out4_1, out4_2, out1], 1) 184 saved_for_loss.append(out4_1) 185 saved_for_loss.append(out4_2) 186 187 out5_1 = self.model5_1(out5) 188 out5_2 = self.model5_2(out5) 189 out6 = torch.cat([out5_1, out5_2, out1], 1) 190 saved_for_loss.append(out5_1) 191 saved_for_loss.append(out5_2) 192 193 out6_1 = self.model6_1(out6) 194 out6_2 = self.model6_2(out6) 195 saved_for_loss.append(out6_1) 196 saved_for_loss.append(out6_2) 197 #其中out6_1 为38个feature maps 198 #其中out6_2 为19个feature maps 199 #saved_for_loss表示需要计算loss的层，对于训练的时候有用 200 return (out6_1, out6_2), saved_for_loss 201 202 def _initialize_weights_norm(self): 203 204 for m in self.modules(): 205 if isinstance(m, nn.Conv2d): 206 init.normal_(m.weight, std=0.01) 207 if m.bias is not None: # mobilenet conv2d doesn't add bias 208 init.constant_(m.bias, 0.0) 209 210 # last layer of these block don't have Relu 211 init.normal_(self.model1_1[8].weight, std=0.01) 212 init.normal_(self.model1_2[8].weight, std=0.01) 213 214 init.normal_(self.model2_1[12].weight, std=0.01) 215 init.normal_(self.model3_1[12].weight, std=0.01) 216 init.normal_(self.model4_1[12].weight, std=0.01) 217 init.normal_(self.model5_1[12].weight, std=0.01) 218 init.normal_(self.model6_1[12].weight, std=0.01) 219 220 init.normal_(self.model2_2[12].weight, std=0.01) 221 init.normal_(self.model3_2[12].weight, std=0.01) 222 init.normal_(self.model4_2[12].weight, std=0.01) 223 init.normal_(self.model5_2[12].weight, std=0.01) 224 init.normal_(self.model6_2[12].weight, std=0.01) 225 226 model = rtpose_model(models) 227 return model 228 229 230 """Load pretrained model on Imagenet 231 :param model, the PyTorch nn.Module which will train. 232 :param model_path, the directory which load the pretrained model, will download one if not have. 233 :param trunk, the feature extractor network of model. 234 """ 235 236 237 def use_vgg(model): 238 239 url = 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth' 240 vgg_state_dict = model_zoo.load_url(url) 241 vgg_keys = vgg_state_dict.keys() 242 243 # load weights of vgg 244 weights_load = {} 245 # weight+bias,weight+bias.....(repeat 10 times) 246 for i in range(20): 247 weights_load[list(model.state_dict().keys())[i] 248 ] = vgg_state_dict[list(vgg_keys)[i]] 249 250 state = model.state_dict() 251 state.update(weights_load) 252 model.load_state_dict(state) 253 print('load imagenet pretrained model') 254
2. get_outputs函数中，包括了模型的搭建，以及前向的工作，得到了paf以及heatmap图。咱们需要特别看下这个函数。
def get_outputs(img, model, preprocess): """Computes the averaged heatmap and paf for the given image :param multiplier: :param origImg: numpy array, the image being processed :param model: pytorch model :returns: numpy arrays, the averaged paf and heatmap """ inp_size = cfg.DATASET.IMAGE_SIZE #其中inp_size为368 # padding #其中的DOWNSAMPLE为默认值是8 im_croped, im_scale, real_shape = im_transform.crop_with_factor( img, inp_size, factor=cfg.MODEL.DOWNSAMPLE, is_ceil=True) #进行图片的处理 #im_cropped size is => (*, 368, 3) if preprocess == 'rtpose': im_data = rtpose_preprocess(im_croped) elif preprocess == 'vgg': im_data = vgg_preprocess(im_croped) elif preprocess == 'inception': im_data = inception_preprocess(im_croped) elif preprocess == 'ssd': im_data = ssd_preprocess(im_croped) batch_images= np.expand_dims(im_data, 0) # several scales as a batch batch_var = torch.from_numpy(batch_images).cuda().float() #其中predicted_outputs是个tuple，后面是用来计算的loss我们用_来接住，不管它 predicted_outputs, _ = model(batch_var) output1, output2 = predicted_outputs[-2], predicted_outputs[-1] heatmap = output2.cpu().data.numpy().transpose(0, 2, 3, 1)[0] paf = output1.cpu().data.numpy().transpose(0, 2, 3, 1)[0] #其中经过了8倍的下采样，(h // 8, w // 8, 38) #(h.//8, w//8, 19) return paf, heatmap, im_scale
3. paf_to_pose_cpp函数，相对而言有点炸，里面涉及的是一个cpp代码，通过swig来进行的。大师总体的工作，，可以通过另一个代码https://blog.csdn.net/l297969586/article/details/80346254来进行对齐，后续有时间我会补上这个swig的代码，相对而言，就是通过一个采样，计算了关键点之间亲和度的方法以及关键点聚类的操作，但是里面有一些值不是很明白，还是说纯粹工程上作者们试出来的，whatever，通过这个函数，我们可以得到图片中有多少个人，以及这些人的关键点坐标，哪些关键点组成哪个人的哪个肢干等等数据，想要的数据都得到了。
4. draw_humans 函数，进行画图，对于上面得到的结果，可以直接拿出每个人，每个人的关节点的坐标，直接进行画点，连线。。
