577 lines
21 KiB
Python
577 lines
21 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
#
|
|
#
|
|
# NOTE: Turn on Internet and GPU
|
|
|
|
# The code hidden below handles all the imports and function definitions (the heavy lifting). If you're a beginner I'd advice you skip this for now. When you are able to understand the rest of the code, come back here and understand each function to get a deeper knowledge.
|
|
|
|
# In[1]:
|
|
|
|
|
|
# !/usr/bin/env python3
|
|
# coding=utf-8
|
|
# author=dave.fang@outlook.com
|
|
# create=20171225
|
|
|
|
import os
|
|
import pprint
|
|
import cv2
|
|
import sys
|
|
import math
|
|
import time
|
|
import tempfile
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.parallel
|
|
import torch.backends.cudnn as cudnn
|
|
import torch.optim as optim
|
|
import torchvision.transforms as transforms
|
|
import torchvision.datasets as datasets
|
|
import torchvision.models as models
|
|
|
|
from torch.autograd import Variable
|
|
|
|
from scipy.ndimage.filters import gaussian_filter
|
|
|
|
#get_ipython().run_line_magic('matplotlib', 'inline')
|
|
#get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")
|
|
|
|
# find connection in the specified sequence, center 29 is in the position 15
|
|
limb_seq = [[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
|
|
map_ids = [[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]]
|
|
|
|
# these are the colours for the 18 body points
|
|
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]]
|
|
|
|
|
|
class PoseEstimation(nn.Module):
|
|
def __init__(self, model_dict):
|
|
super(PoseEstimation, self).__init__()
|
|
|
|
self.model0 = model_dict['block_0']
|
|
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
|
|
|
|
|
|
def make_layers(layer_dict):
|
|
layers = []
|
|
|
|
for i in range(len(layer_dict) - 1):
|
|
layer = layer_dict[i]
|
|
for k in layer:
|
|
v = layer[k]
|
|
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)]
|
|
layer = list(layer_dict[-1].keys())
|
|
k = layer[0]
|
|
v = layer_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)
|
|
|
|
|
|
def get_pose_model():
|
|
blocks = {}
|
|
|
|
block_0 = [{'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]}]
|
|
|
|
layers = []
|
|
for block in block_0:
|
|
# print(block)
|
|
for key in block:
|
|
v = block[key]
|
|
if 'pool' in key:
|
|
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 = {
|
|
'block_0': nn.Sequential(*layers)
|
|
}
|
|
|
|
for k in blocks:
|
|
v = blocks[k]
|
|
models[k] = make_layers(v)
|
|
|
|
return PoseEstimation(models)
|
|
|
|
|
|
def get_paf_and_heatmap(model, img_raw, scale_search, param_stride=8, box_size=368):
|
|
multiplier = [scale * box_size / img_raw.shape[0] for scale in scale_search]
|
|
|
|
heatmap_avg = torch.zeros((len(multiplier), 19, img_raw.shape[0], img_raw.shape[1])).cuda()
|
|
paf_avg = torch.zeros((len(multiplier), 38, img_raw.shape[0], img_raw.shape[1])).cuda()
|
|
|
|
for i, scale in enumerate(multiplier):
|
|
img_test = cv2.resize(img_raw, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
|
|
img_test_pad, pad = pad_right_down_corner(img_test, param_stride, param_stride)
|
|
img_test_pad = np.transpose(np.float32(img_test_pad[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
|
|
|
|
feed = Variable(torch.from_numpy(img_test_pad)).cuda()
|
|
output1, output2 = model(feed)
|
|
|
|
print(output1.size())
|
|
print(output2.size())
|
|
|
|
heatmap = nn.UpsamplingBilinear2d((img_raw.shape[0], img_raw.shape[1])).cuda()(output2)
|
|
|
|
paf = nn.UpsamplingBilinear2d((img_raw.shape[0], img_raw.shape[1])).cuda()(output1)
|
|
|
|
heatmap_avg[i] = heatmap[0].data
|
|
paf_avg[i] = paf[0].data
|
|
|
|
heatmap_avg = torch.transpose(torch.transpose(torch.squeeze(torch.mean(heatmap_avg, 0)), 0, 1), 1, 2).cuda()
|
|
heatmap_avg = heatmap_avg.cpu().numpy()
|
|
|
|
paf_avg = torch.transpose(torch.transpose(torch.squeeze(torch.mean(paf_avg, 0)), 0, 1), 1, 2).cuda()
|
|
paf_avg = paf_avg.cpu().numpy()
|
|
|
|
return paf_avg, heatmap_avg
|
|
|
|
|
|
def extract_heatmap_info(heatmap_avg, param_thre1=0.1):
|
|
all_peaks = []
|
|
peak_counter = 0
|
|
|
|
for part in range(18):
|
|
map_ori = heatmap_avg[:, :, part]
|
|
map_gau = gaussian_filter(map_ori, sigma=3)
|
|
|
|
map_left = np.zeros(map_gau.shape)
|
|
map_left[1:, :] = map_gau[:-1, :]
|
|
map_right = np.zeros(map_gau.shape)
|
|
map_right[:-1, :] = map_gau[1:, :]
|
|
map_up = np.zeros(map_gau.shape)
|
|
map_up[:, 1:] = map_gau[:, :-1]
|
|
map_down = np.zeros(map_gau.shape)
|
|
map_down[:, :-1] = map_gau[:, 1:]
|
|
|
|
peaks_binary = np.logical_and.reduce(
|
|
(map_gau >= map_left, map_gau >= map_right, map_gau >= map_up,
|
|
map_gau >= map_down, map_gau > param_thre1))
|
|
|
|
peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
|
|
peaks = list(peaks)
|
|
peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
|
|
ids = range(peak_counter, peak_counter + len(peaks))
|
|
peaks_with_score_and_id = [peaks_with_score[i] + (ids[i],) for i in range(len(ids))]
|
|
|
|
all_peaks.append(peaks_with_score_and_id)
|
|
peak_counter += len(peaks)
|
|
|
|
return all_peaks
|
|
|
|
|
|
def extract_paf_info(img_raw, paf_avg, all_peaks, param_thre2=0.05, param_thre3=0.5):
|
|
connection_all = []
|
|
special_k = []
|
|
mid_num = 10
|
|
|
|
for k in range(len(map_ids)):
|
|
score_mid = paf_avg[:, :, [x - 19 for x in map_ids[k]]]
|
|
candA = all_peaks[limb_seq[k][0] - 1]
|
|
candB = all_peaks[limb_seq[k][1] - 1]
|
|
nA = len(candA)
|
|
nB = len(candB)
|
|
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))
|
|
startend = list(startend)
|
|
|
|
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))])
|
|
|
|
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
|
|
score_with_dist_prior = sum(score_midpts) / len(score_midpts)
|
|
score_with_dist_prior += min(0.5 * img_raw.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:
|
|
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)
|
|
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]])
|
|
if len(connection) >= min(nA, nB):
|
|
break
|
|
|
|
connection_all.append(connection)
|
|
else:
|
|
special_k.append(k)
|
|
connection_all.append([])
|
|
|
|
return special_k, connection_all
|
|
|
|
|
|
def get_subsets(connection_all, special_k, all_peaks):
|
|
# last number in each row is the total parts number of that person
|
|
# the second last number in each row is the score of the overall configuration
|
|
subset = -1 * np.ones((0, 20))
|
|
candidate = np.array([item for sublist in all_peaks for item in sublist])
|
|
|
|
for k in range(len(map_ids)):
|
|
if k not in special_k:
|
|
partAs = connection_all[k][:, 0]
|
|
partBs = connection_all[k][:, 1]
|
|
indexA, indexB = np.array(limb_seq[k]) - 1
|
|
|
|
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):
|
|
if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
|
|
subset_idx[found] = j
|
|
found += 1
|
|
|
|
if found == 1:
|
|
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 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: # merge
|
|
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)
|
|
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])
|
|
return subset, candidate
|
|
|
|
|
|
def draw_key_point(subset, all_peaks, img_raw):
|
|
del_ids = []
|
|
for i in range(len(subset)):
|
|
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
|
|
del_ids.append(i)
|
|
subset = np.delete(subset, del_ids, axis=0)
|
|
|
|
img_canvas = img_raw.copy() # B,G,R order
|
|
|
|
for i in range(18):
|
|
for j in range(len(all_peaks[i])):
|
|
cv2.circle(img_canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
|
|
|
|
return subset, img_canvas
|
|
|
|
|
|
def link_key_point(img_canvas, candidate, subset, stickwidth=4):
|
|
for i in range(17):
|
|
for n in range(len(subset)):
|
|
index = subset[n][np.array(limb_seq[i]) - 1]
|
|
if -1 in index:
|
|
continue
|
|
cur_canvas = img_canvas.copy()
|
|
Y = candidate[index.astype(int), 0]
|
|
X = candidate[index.astype(int), 1]
|
|
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])
|
|
img_canvas = cv2.addWeighted(img_canvas, 0.4, cur_canvas, 0.6, 0)
|
|
|
|
return img_canvas
|
|
|
|
def pad_right_down_corner(img, stride, pad_value):
|
|
h = img.shape[0]
|
|
w = img.shape[1]
|
|
|
|
pad = 4 * [None]
|
|
pad[0] = 0 # up
|
|
pad[1] = 0 # left
|
|
pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
|
|
pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
|
|
|
|
img_padded = img
|
|
pad_up = np.tile(img_padded[0:1, :, :] * 0 + pad_value, (pad[0], 1, 1))
|
|
img_padded = np.concatenate((pad_up, img_padded), axis=0)
|
|
pad_left = np.tile(img_padded[:, 0:1, :] * 0 + pad_value, (1, pad[1], 1))
|
|
img_padded = np.concatenate((pad_left, img_padded), axis=1)
|
|
pad_down = np.tile(img_padded[-2:-1, :, :] * 0 + pad_value, (pad[2], 1, 1))
|
|
img_padded = np.concatenate((img_padded, pad_down), axis=0)
|
|
pad_right = np.tile(img_padded[:, -2:-1, :] * 0 + pad_value, (1, pad[3], 1))
|
|
img_padded = np.concatenate((img_padded, pad_right), axis=1)
|
|
|
|
return img_padded, pad
|
|
|
|
|
|
if __name__ == '__main__':
|
|
print(get_pose_model())
|
|
|
|
|
|
# First let's download the pre-trained model.
|
|
|
|
# In[2]:
|
|
|
|
|
|
# Using gdown to download the model directly from Google Drive
|
|
|
|
#assert os.system(' conda install -y gdown') == 0
|
|
import gdown
|
|
|
|
|
|
# In[3]:
|
|
|
|
|
|
model = 'coco_pose_iter_440000.pth.tar'
|
|
if not os.path.exists(model):
|
|
url = 'https://drive.google.com/u/0/uc?export=download&confirm=f_Ix&id=0B1asvDK18cu_MmY1ZkpaOUhhRHM'
|
|
gdown.download(
|
|
url,
|
|
model,
|
|
quiet=False
|
|
)
|
|
|
|
|
|
# In[4]:
|
|
|
|
|
|
state_dict = torch.load('./coco_pose_iter_440000.pth.tar')['state_dict'] # getting the pre-trained model's parameters
|
|
# A state_dict is simply a Python dictionary object that maps each layer to its parameter tensor.
|
|
|
|
model_pose = get_pose_model() # building the model (see fn. defn. above). To see the architecture, see below cell.
|
|
model_pose.load_state_dict(state_dict) # Loading the parameters (weights, biases) into the model.
|
|
|
|
model_pose.float() # I'm not sure why this is used. No difference if you remove it.
|
|
|
|
|
|
# In[5]:
|
|
|
|
|
|
arch_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0000.jpg'
|
|
img_ori = cv2.imread(arch_image)
|
|
plt.figure(figsize=(15, 8))
|
|
plt.imshow(img_ori[...,::-1])
|
|
|
|
|
|
# Notice, the first 10 layers are from VGG-19. But here instead of downloading the model and loading the layers from there, we simply hardcoaded it in get_pose_model()
|
|
|
|
# In[6]:
|
|
|
|
|
|
# Run this to view the model's architecture
|
|
#model_pose.eval()
|
|
|
|
|
|
# In[7]:
|
|
|
|
|
|
use_gpu = True
|
|
|
|
if use_gpu:
|
|
model_pose.cuda()
|
|
model_pose = torch.nn.DataParallel(model_pose, device_ids=range(torch.cuda.device_count()))
|
|
cudnn.benchmark = True
|
|
|
|
|
|
# In[8]:
|
|
|
|
|
|
def estimate_pose(img_ori, name=None):
|
|
if name is None:
|
|
name = tempfile.mktemp(
|
|
dir='/kaggle/working',
|
|
suffix='.png',
|
|
)
|
|
pprint.pprint(
|
|
['estimate_pose', dict(name=name)],
|
|
)
|
|
|
|
# People might be at different scales in the image, perform inference at multiple scales to boost results
|
|
scale_param = [0.5, 1.0, 1.5, 2.0]
|
|
|
|
# Predict Heatmaps for approximate joint position
|
|
# Use Part Affinity Fields (PAF's) as guidance to link joints to form skeleton
|
|
# PAF's are just unit vectors along the limb encoding the direction of the limb
|
|
# A dot product of possible joint connection will be high if actual limb else low
|
|
|
|
paf_info, heatmap_info = get_paf_and_heatmap(model_pose, img_ori, scale_param)
|
|
peaks = extract_heatmap_info(heatmap_info)
|
|
sp_k, con_all = extract_paf_info(img_ori, paf_info, peaks)
|
|
|
|
subsets, candidates = get_subsets(con_all, sp_k, peaks)
|
|
subsets, img_points = draw_key_point(subsets, peaks, img_ori)
|
|
|
|
# After predicting Heatmaps and PAF's, proceeed to link joints correctly
|
|
img_canvas = link_key_point(img_points, candidates, subsets)
|
|
|
|
|
|
f = plt.figure(figsize=(15, 10))
|
|
|
|
plt.subplot(1, 2, 1)
|
|
plt.imshow(img_points[...,::-1])
|
|
|
|
plt.subplot(1, 2, 2)
|
|
plt.imshow(img_canvas[...,::-1])
|
|
|
|
f.savefig(name)
|
|
|
|
|
|
# In[9]:
|
|
|
|
|
|
test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0000.jpg'
|
|
img_ori = cv2.imread(test_image)
|
|
estimate_pose(img_ori)
|
|
|
|
|
|
# In[10]:
|
|
|
|
|
|
test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0010.jpg'
|
|
img_ori = cv2.imread(test_image)
|
|
estimate_pose(img_ori)
|
|
|
|
|
|
# In[11]:
|
|
|
|
|
|
test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0020.jpg'
|
|
img_ori = cv2.imread(test_image)
|
|
estimate_pose(img_ori)
|
|
|
|
|
|
# In[12]:
|
|
|
|
|
|
test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0030.jpg'
|
|
img_ori = cv2.imread(test_image)
|
|
estimate_pose(img_ori)
|
|
|
|
|
|
# In[13]:
|
|
|
|
|
|
test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0040.jpg'
|
|
img_ori = cv2.imread(test_image)
|
|
estimate_pose(img_ori)
|
|
|
|
|
|
# In[14]:
|
|
|
|
|
|
test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0050.jpg'
|
|
img_ori = cv2.imread(test_image)
|
|
estimate_pose(img_ori)
|
|
|
|
|
|
# In[ ]:
|