dummy_pose_estimation.py

import numpy as np
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from scipy.misc import imresize, imread

from human_pose_nn import HumanPoseIRNetwork

net_pose = HumanPoseIRNetwork()
net_pose.restore('models/MPII+LSP.ckpt')

img = imread('images/dummy.jpg')
img = imresize(img, [299, 299])
img_batch = np.expand_dims(img, 0)

y, x, a = net_pose.estimate_joints(img_batch)
y, x, a = np.squeeze(y), np.squeeze(x), np.squeeze(a)

joint_names = [
    'right ankle ',
    'right knee ',
    'right hip',
    'left hip',
    'left knee',
    'left ankle',
    'pelvis',
    'thorax',
    'upper neck',
    'head top',
    'right wrist',
    'right elbow',
    'right shoulder',
    'left shoulder',
    'left elbow',
    'left wrist'
]

# Print probabilities of each estimation
for i in range(16):
    print('%s: %.02f%%' % (joint_names[i], a[i] * 100))

# Create image
colors = ['r', 'r', 'b', 'm', 'm', 'y', 'g', 'g', 'b', 'c', 'r', 'r', 'b', 'm', 'm', 'c']
for i in range(16):
    if i < 15 and i not in {5, 9}:
        plt.plot([x[i], x[i + 1]], [y[i], y[i + 1]], color = colors[i], linewidth = 5)

plt.imshow(img)
plt.savefig('images/dummy_pose.jpg')