#!/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[ ]: