github: https://github.com/tensorboy/pytorch_Realtime_Multi-Person_Pose_Estimation

 # -*- coding: utf-8 -*

 import os

 import re

 import sys

 import cv2

 import math

 import time

 import scipy

 import argparse

 import matplotlib

 import numpy as np

 import pylab as plt

 from joblib import Parallel, delayed

 import util

 import torch

 import torch as T

 import torch.nn as nn

 import torch.nn.functional as F

 from torch.autograd import Variable

 from collections import OrderedDict

 from config_reader import config_reader

 from scipy.ndimage.filters import gaussian_filter

 #parser = argparse.ArgumentParser()

 #parser.add_argument('--t7_file', required=True)

 #parser.add_argument('--pth_file', required=True)

 #args = parser.parse_args()

 torch.set_num_threads(torch.get_num_threads())

 weight_name = './model/pose_model.pth'

 blocks = {}

 # 从1开始算的limb,图对应:Pose Output Format

 # find connection in the specified sequence, center 29 is in the position 15

 limbSeq = [[2,3], [2,6], [3,4], [4,5], [6,7], [7,8], [2,9], [9,10], \

            [10,11], [2,12], [12,13], [13,14], [2,1], [1,15], [15,17], \

            [1,16], [16,18], [3,17], [6,18]]

 # the middle joints heatmap correpondence

 mapIdx = [[31,32], [39,40], [33,34], [35,36], [41,42], [43,44], [19,20], [21,22], \

           [23,24], [25,26], [27,28], [29,30], [47,48], [49,50], [53,54], [51,52], \

           [55,56], [37,38], [45,46]]

 # visualize

 colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \

           [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \

           [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]

 # heatmap channel为19 表示关节点的score

 # PAF channel为38 表示limb的单位向量

 block0  = [{'conv1_1':[3,64,3,1,1]},{'conv1_2':[64,64,3,1,1]},{'pool1_stage1':[2,2,0]},{'conv2_1':[64,128,3,1,1]},{'conv2_2':[128,128,3,1,1]},{'pool2_stage1':[2,2,0]},{'conv3_1':[128,256,3,1,1]},{'conv3_2':[256,256,3,1,1]},{'conv3_3':[256,256,3,1,1]},{'conv3_4':[256,256,3,1,1]},{'pool3_stage1':[2,2,0]},{'conv4_1':[256,512,3,1,1]},{'conv4_2':[512,512,3,1,1]},{'conv4_3_CPM':[512,256,3,1,1]},{'conv4_4_CPM':[256,128,3,1,1]}]

 blocks['block1_1']  = [{'conv5_1_CPM_L1':[128,128,3,1,1]},{'conv5_2_CPM_L1':[128,128,3,1,1]},{'conv5_3_CPM_L1':[128,128,3,1,1]},{'conv5_4_CPM_L1':[128,512,1,1,0]},{'conv5_5_CPM_L1':[512,38,1,1,0]}]

 blocks['block1_2']  = [{'conv5_1_CPM_L2':[128,128,3,1,1]},{'conv5_2_CPM_L2':[128,128,3,1,1]},{'conv5_3_CPM_L2':[128,128,3,1,1]},{'conv5_4_CPM_L2':[128,512,1,1,0]},{'conv5_5_CPM_L2':[512,19,1,1,0]}]

 for i in range(2,7):

     blocks['block%d_1'%i]  = [{'Mconv1_stage%d_L1'%i:[185,128,7,1,3]},{'Mconv2_stage%d_L1'%i:[128,128,7,1,3]},{'Mconv3_stage%d_L1'%i:[128,128,7,1,3]},{'Mconv4_stage%d_L1'%i:[128,128,7,1,3]},

 {'Mconv5_stage%d_L1'%i:[128,128,7,1,3]},{'Mconv6_stage%d_L1'%i:[128,128,1,1,0]},{'Mconv7_stage%d_L1'%i:[128,38,1,1,0]}]

     blocks['block%d_2'%i]  = [{'Mconv1_stage%d_L2'%i:[185,128,7,1,3]},{'Mconv2_stage%d_L2'%i:[128,128,7,1,3]},{'Mconv3_stage%d_L2'%i:[128,128,7,1,3]},{'Mconv4_stage%d_L2'%i:[128,128,7,1,3]},

 {'Mconv5_stage%d_L2'%i:[128,128,7,1,3]},{'Mconv6_stage%d_L2'%i:[128,128,1,1,0]},{'Mconv7_stage%d_L2'%i:[128,19,1,1,0]}]

 def make_layers(cfg_dict):

     layers = []

     for i in range(len(cfg_dict)-1):

         one_ = cfg_dict[i]

         for k,v in one_.iteritems():

             if 'pool' in k:

                 layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2] )]

             else:

                 conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride = v[3], padding=v[4])

                 layers += [conv2d, nn.ReLU(inplace=True)]

     one_ = cfg_dict[-1].keys()

     k = one_[0]

     v = cfg_dict[-1][k]

     conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride = v[3], padding=v[4])

     layers += [conv2d]

     return nn.Sequential(*layers)

 layers = []

 for i in range(len(block0)):

     one_ = block0[i]

     for k,v in one_.iteritems():

         if 'pool' in k:

             layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2] )]

         else:

             conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride = v[3], padding=v[4])

             layers += [conv2d, nn.ReLU(inplace=True)]  

 models = {}

 models['block0']=nn.Sequential(*layers)        

 for k,v in blocks.iteritems():

     models[k] = make_layers(v)

 class pose_model(nn.Module):

     def __init__(self,model_dict,transform_input=False):

         super(pose_model, self).__init__()

         self.model0   = model_dict['block0']

         self.model1_1 = model_dict['block1_1']

         self.model2_1 = model_dict['block2_1']

         self.model3_1 = model_dict['block3_1']

         self.model4_1 = model_dict['block4_1']

         self.model5_1 = model_dict['block5_1']

         self.model6_1 = model_dict['block6_1']  

         self.model1_2 = model_dict['block1_2']

         self.model2_2 = model_dict['block2_2']

         self.model3_2 = model_dict['block3_2']

         self.model4_2 = model_dict['block4_2']

         self.model5_2 = model_dict['block5_2']

         self.model6_2 = model_dict['block6_2']

     def forward(self, x):

         out1 = self.model0(x)

         out1_1 = self.model1_1(out1)

         out1_2 = self.model1_2(out1)

         out2  = torch.cat([out1_1,out1_2,out1],1)

         out2_1 = self.model2_1(out2)

         out2_2 = self.model2_2(out2)

         out3   = torch.cat([out2_1,out2_2,out1],1)

         out3_1 = self.model3_1(out3)

         out3_2 = self.model3_2(out3)

         out4   = torch.cat([out3_1,out3_2,out1],1)

         out4_1 = self.model4_1(out4)

         out4_2 = self.model4_2(out4)

         out5   = torch.cat([out4_1,out4_2,out1],1)  

         out5_1 = self.model5_1(out5)

         out5_2 = self.model5_2(out5)

         out6   = torch.cat([out5_1,out5_2,out1],1)         

         out6_1 = self.model6_1(out6)

         out6_2 = self.model6_2(out6)

         return out6_1,out6_2        

 model = pose_model(models)

 model.load_state_dict(torch.load(weight_name))

 model.cuda()

 model.float()

 model.eval()

 param_, model_ = config_reader()

 #torch.nn.functional.pad(img pad, mode='constant', value=model_['padValue'])

 tic = time.time()

 test_image = './sample_image/ski.jpg'

 #test_image = 'a.jpg'

 oriImg = cv2.imread(test_image) # B,G,R order

 imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(),0),2,3),1,2),volatile=True).cuda()

 multiplier = [x * model_['boxsize'] / oriImg.shape[0] for x in param_['scale_search']] # 不同scale输入

 heatmap_avg = torch.zeros((len(multiplier),19,oriImg.shape[0], oriImg.shape[1])).cuda()

 paf_avg = torch.zeros((len(multiplier),38,oriImg.shape[0], oriImg.shape[1])).cuda()

 #print heatmap_avg.size()

 toc =time.time()

 print 'time is %.5f'%(toc-tic)

 tic = time.time()

 for m in range(len(multiplier)):

     scale = multiplier[m]

     h = int(oriImg.shape[0]*scale)

     w = int(oriImg.shape[1]*scale)

     pad_h = 0 if (h%model_['stride']==0) else model_['stride'] - (h % model_['stride'])

     pad_w = 0 if (w%model_['stride']==0) else model_['stride'] - (w % model_['stride'])

     new_h = h+pad_h

     new_w = w+pad_w

     imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)

     imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_['stride'], model_['padValue'])

     imageToTest_padded = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5

     # (-0.5~0.5)

     feed = Variable(T.from_numpy(imageToTest_padded)).cuda()

     output1,output2 = model(feed)

     print output1.size()

     print output2.size()

     heatmap = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output2) # 对output上采样至原图大小

     paf = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output1)     # 同理

     heatmap_avg[m] = heatmap[0].data

     paf_avg[m] = paf[0].data  

 toc =time.time()

 print 'time is %.5f'%(toc-tic)

 tic = time.time()

 # 不同scale的heatmap和PAF取均值

 heatmap_avg = T.transpose(T.transpose(T.squeeze(T.mean(heatmap_avg, 0)),0,1),1,2).cuda()

 paf_avg     = T.transpose(T.transpose(T.squeeze(T.mean(paf_avg, 0)),0,1),1,2).cuda()

 heatmap_avg=heatmap_avg.cpu().numpy()

 paf_avg    = paf_avg.cpu().numpy()

 toc =time.time()

 print 'time is %.5f'%(toc-tic)

 tic = time.time()

 all_peaks = []

 peak_counter = 0

 #maps =

 # picture array is reversed

 for part in range(18): # 18个关节点的featuremap

     map_ori = heatmap_avg[:,:,part]

     map = gaussian_filter(map_ori, sigma=3)

     map_left = np.zeros(map.shape)

     map_left[1:,:] = map[:-1,:]

     map_right = np.zeros(map.shape)

     map_right[:-1,:] = map[1:,:]

     map_up = np.zeros(map.shape)

     map_up[:,1:] = map[:,:-1]

     map_down = np.zeros(map.shape)

     map_down[:,:-1] = map[:,1:]

     # 计算是否为局部极值

     peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))

 #    peaks_binary = T.eq(

 #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])

     peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse

     peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]

     id = range(peak_counter, peak_counter + len(peaks))

     peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]

     all_peaks.append(peaks_with_score_and_id) # 一个关节点featuremap上不同人的peak [[y, x, peak_score, id)],...]

     peak_counter += len(peaks)

 # 计算线性积分 采样10个点计算

 connection_all = []

 special_k = []

 mid_num = 10

 for k in range(len(mapIdx)):

     score_mid = paf_avg[:,:,[x-19 for x in mapIdx[k]]] # channel为2的paf_avg,表示PAF向量

     candA = all_peaks[limbSeq[k][0]-1]  #第k个limb中左关节点的候选集合A(不同人的关节点)

     candB = all_peaks[limbSeq[k][1]-1]  #第k个limb中右关节点的候选集合B(不同人的关节点)

     nA = len(candA)

     nB = len(candB)

     # indexA, indexB = limbSeq[k]

     if(nA != 0 and nB != 0): # 有候选时开始连接

         connection_candidate = []

         for i in range(nA):

             for j in range(nB):

                 vec = np.subtract(candB[j][:2], candA[i][:2])

                 norm = math.sqrt(vec[0]*vec[0] + vec[1]*vec[1])

                 vec = np.divide(vec, norm)  # 计算单位向量

                 startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \

                                np.linspace(candA[i][1], candB[j][1], num=mid_num)) # 在A[i],B[j]连接线上采样mid_num个点

                 # 计算采样点的PAF向量

                 vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \

                                   for I in range(len(startend))])

                 vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \

                                   for I in range(len(startend))])

                 # 采样点的PAF向量与limb的单位向量计算余弦相似度score,内积

                 score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])

                 score_with_dist_prior = sum(score_midpts)/len(score_midpts) + min(0.5*oriImg.shape[0]/norm-1, 0)

                 criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)

                 criterion2 = score_with_dist_prior > 0

                 if criterion1 and criterion2:

                     # (i,j,score,score_all)

                     connection_candidate.append([i, j, score_with_dist_prior, score_with_dist_prior+candA[i][2]+candB[j][2]])

         connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True) # 按score排序

         connection = np.zeros((0,5))

         for c in range(len(connection_candidate)):

             i,j,s = connection_candidate[c][0:3]

             if(i not in connection[:,3] and j not in connection[:,4]):

                 connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]]) # A_id, B_id, score, i, j

                 if(len(connection) >= min(nA, nB)):

                     break

         connection_all.append(connection) # 多个符合当前limb的组合 [[A_id, B_id, score, i, j],...]

     else:

         special_k.append(k)

         connection_all.append([])

 '''

 function: 关节点连成每个人的limb

 subset: last number in each row is the total parts number of that person

 subset: the second last number in each row is the score of the overall configuration

 candidate: 候选关节点

 connection_all: 候选limb

 '''

 subset = -1 * np.ones((0, 20))

 candidate = np.array([item for sublist in all_peaks for item in sublist])  # 一个id的(y,x,score,id)(关节点)

 for k in range(len(mapIdx)):

     if k not in special_k:

         partAs = connection_all[k][:,0]  # 第k个limb,左端点的候选id集合

         partBs = connection_all[k][:,1]  # 第k个limb,右端点的候选id集合

         indexA, indexB = np.array(limbSeq[k]) - 1 # 关节点index

         for i in range(len(connection_all[k])): #= 1:size(temp,1)

             found = 0

             subset_idx = [-1, -1]

             for j in range(len(subset)): #1:size(subset,1): 遍历subset里每个人,看当前两个关节点出现过几次

                 if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:

                     subset_idx[found] = j

                     found += 1

             if found == 1: # 在这个人的subset里连上这个limb

                 j = subset_idx[0]

                 if(subset[j][indexB] != partBs[i]):

                     subset[j][indexB] = partBs[i]

                     subset[j][-1] += 1

                     subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]

                 elif(subset[j][indexA] != partAs[i]):

                     subset[j][indexA] = partAs[i]

                     subset[j][-1] += 1

                     subset[j][-2] += candidate[partAs[i].astype(int), 2] + connection_all[k][i][2]

             elif found == 2: # if found 2 and disjoint, merge them

                 j1, j2 = subset_idx

                 print "found = 2"

                 membership = ((subset[j1]>=0).astype(int) + (subset[j2]>=0).astype(int))[:-2]

                 if len(np.nonzero(membership == 2)[0]) == 0:

                     # 如果两个人的相同关节点没有在各自的subset中都连成limb,那么合并两个subset构成一个人

                     subset[j1][:-2] += (subset[j2][:-2] + 1)

                     subset[j1][-2:] += subset[j2][-2:]

                     subset[j1][-2] += connection_all[k][i][2]

                     subset = np.delete(subset, j2, 0)

                 # To-Do 这里有问题, 怎么合并才对?

                 # else: # as like found == 1

                 #     subset[j1][indexB] = partBs[i]

                 #     subset[j1][-1] += 1

                 #     subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]

             # if find no partA in the subset, create a new subset

             elif not found and k < 17:

                 row = -1 * np.ones(20)

                 row[indexA] = partAs[i]

                 row[indexB] = partBs[i]

                 row[-1] = 2

                 row[-2] = sum(candidate[connection_all[k][i,:2].astype(int), 2]) + connection_all[k][i][2]

                 subset = np.vstack([subset, row])

 # delete some rows of subset which has few parts occur

 deleteIdx = [];

 for i in range(len(subset)):

     if subset[i][-1] < 4 or subset[i][-2]/subset[i][-1] < 0.4:

         deleteIdx.append(i)

 subset = np.delete(subset, deleteIdx, axis=0)

 canvas = cv2.imread(test_image) # B,G,R order

 for i in range(18):

     for j in range(len(all_peaks[i])):

         cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)

 stickwidth = 4

 for i in range(17):

     for n in range(len(subset)):

         index = subset[n][np.array(limbSeq[i])-1] # limb的两个关节点index

         if -1 in index:

             continue

         cur_canvas = canvas.copy()

         Y = candidate[index.astype(int), 0] # 两个index点的纵坐标

         X = candidate[index.astype(int), 1] # 两个index点的横坐标

         mX = np.mean(X)

         mY = np.mean(Y)

         length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5

         angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))

         polygon = cv2.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)

         cv2.fillConvexPoly(cur_canvas, polygon, colors[i])

         canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)

 #Parallel(n_jobs=1)(delayed(handle_one)(i) for i in range(18))

 toc =time.time()

 print 'time is %.5f'%(toc-tic)

 cv2.imwrite('result.png',canvas)

openpose pytorch代码分析的更多相关文章

  1. （原）SphereFace及其pytorch代码

    转载请注明出处: http://www.cnblogs.com/darkknightzh/p/8524937.html 论文: SphereFace: Deep Hypersphere Embeddi ...

  2. 目标检测之Faster-RCNN的pytorch代码详解(数据预处理篇)

    首先贴上代码原作者的github:https://github.com/chenyuntc/simple-faster-rcnn-pytorch(非代码作者,博文只解释代码) 今天看完了simple- ...

  3. 残差网络resnet理解与pytorch代码实现

    写在前面 ​ 深度残差网络(Deep residual network, ResNet)自提出起,一次次刷新CNN模型在ImageNet中的成绩,解决了CNN模型难训练的问题.何凯明大神的工作令人佩服 ...

  4. Android代码分析工具lint学习

    1 lint简介 1.1 概述 lint是随Android SDK自带的一个静态代码分析工具.它用来对Android工程的源文件进行检查,找出在正确性.安全.性能.可使用性.可访问性及国际化等方面可能 ...

  5. pmd静态代码分析

    在正式进入测试之前,进行一定的静态代码分析及code review对代码质量及系统提高是有帮助的,以上为数据证明 Pmd 它是一个基于静态规则集的Java源码分析器,它可以识别出潜在的如下问题:– 可 ...

  6. [Asp.net 5] DependencyInjection项目代码分析-目录

    微软DI文章系列如下所示: [Asp.net 5] DependencyInjection项目代码分析 [Asp.net 5] DependencyInjection项目代码分析2-Autofac [ ...

  7. [Asp.net 5] DependencyInjection项目代码分析4-微软的实现（5）（IEnumerable<>补充）

    Asp.net 5的依赖注入注入系列可以参考链接: [Asp.net 5] DependencyInjection项目代码分析-目录 我们在之前讲微软的实现时,对于OpenIEnumerableSer ...

  8. 完整全面的Java资源库（包括构建、操作、代码分析、编译器、数据库、社区等等）

    构建 这里搜集了用来构建应用程序的工具. Apache Maven:Maven使用声明进行构建并进行依赖管理,偏向于使用约定而不是配置进行构建.Maven优于Apache Ant.后者采用了一种过程化 ...

  9. STM32启动代码分析 IAR 比较好

    stm32启动代码分析 (2012-06-12 09:43:31) 转载▼     最近开始使用ST的stm32w108芯片(也是一款zigbee芯片).开始看他的启动代码看的晕晕呼呼呼的. 还好在c ...

随机推荐

  1. PCM EQ DRC 音频处理

    PCM Pulse-code modulation的缩写,中文译名是脉冲编码调制.(I2S仅仅是PCM的一个分支,接口定义都是一样的, I2S的采样频率一般为44.1KHZ和48KHZ做,PCM采样频 ...

  2. noi.ac 集合

    A.集合 --- 题面 不知道有没有用的传送门[滑稽 就是给你一个 包含 1~n 的集合,让你求它的大小为 k 的子集 s 的 \(T^{min(s)}\) 的期望值, T 为给出值, min(s) ...

  3. PHP一维数组转二维数组正则表达式

    2017年11月20日17:17:08 array(1 => '哈哈')  变成  array('id' => 1, 'name' => '哈哈') 查找目标:  (\d)\s=&g ...

  4. mysql alter 效率

    2017年9月15日 10:36:54 星期五 今天遇到一个效率问题记下来: 场景: mysql要更改一下表字段的注释, 因为sql语句问题, 导致更新了整张表.. 错误: ) UNSIGNED ' ...

  5. 如何将本地项目上传到Github

    看了这篇文章觉得写的很详细很适合初学者  提供给大家参考下. http://blog.csdn.net/zamamiro/article/details/70172900 注意如果第二次git pus ...

  6. IOS 常遇到的报错警告 以及 解决办法

    1.  This application is modifying the autolayout engine from a background thread, which can lead to ...

  7. VM_Centos7.3_X64_安装Oracle12C 总结笔记

    声明:本文居多内容参考原文来之网络: 一:安装Centos7.3 虚拟机 1:操作系统下载 CentOS 7官方下载地址:https://www.centos.org/download/ 说明:本案例 ...

  8. Confluence 6 管理协同编辑 - 审计的考虑

    我们知道一些客户对审计是主要考虑的方面.我们不能保证在协同编辑的时候具有审计,审查功能.所有页面的修改当前附加到用户发布页面的属性中而不是用户的特定修改. 如果这个对你来说是一个问题的话,我们建议你在 ...

  9. Confluence 6 那些文件需要备份

    备份整个 home 目录是最安全的选项.但是,有很多目录是在 Confluence 启动的时候创建的并且也是可以忽略的.不管那些文件夹可以忽略,下面的文件夹必须进行备份才能回复: <conf-h ...

  10. nginx实践(二)之静态资源web服务(浏览器缓存场景）

    配置语法-expires