def visualize_short_offsets(offsets,
keypoint_id,
centers=None,
heatmaps=None,
radius=config.KP_RADIUS,
img=None,
every=1,
n=0):
if centers is None and heatmaps is None:
raise ValueError(
'either keypoint locations or heatmaps must be provided')
if isinstance(keypoint_id, str):
if not keypoint_id in config.KEYPOINTS:
raise ValueError(
'{} not a valid keypoint name'.format(keypoint_id))
else:
keypoint_id = config.KEYPOINTS.index(keypoint_id)
if centers is None:
kp = get_keypoints(heatmaps)
kp = [k for k in kp if k['id'] == keypoint_id]
centers = [k['xy'].tolist() for k in kp]
kp_offsets = offsets[:, :, 2 * keypoint_id:2 * keypoint_id + 2]
masks = np.zeros(offsets.shape[:2] + (len(centers), ), dtype='bool')
idx = np.rollaxis(np.indices(offsets.shape[1::-1]), 0, 3).transpose(
(1, 0, 2))
for j, c in enumerate(centers):
dists = np.sqrt(np.square(idx - c).sum(axis=-1))
dists_x = np.abs(idx[:, :, 0] - c[0])
dists_y = np.abs(idx[:, :, 1] - c[1])
masks[:, :, j] = (dists <= radius)
if every > 1:
d_mask = np.logical_and(
np.mod(dists_x.astype('int32'), every) == 0,
np.mod(dists_y.astype('int32'), every) == 0)
masks[:, :, j] = np.logical_and(masks[:, :, j], d_mask)
mask = masks.sum(axis=-1) > 0
# for j, c in enumerate(centers):
# dists[:,:,j] = np.sqrt(np.square(idx-c).sum(axis=-1))
# dists = dists.min(axis=-1)
# mask = dists <= radius
I, J = np.nonzero(mask)
plt.figure()
if img is not None:
plt.imshow(img)
plt.quiver(J,
I,
kp_offsets[I, J, 0],
kp_offsets[I, J, 1],
color='r',
angles='xy',
scale_units='xy',
scale=1)
plt.axis('off')
plt.savefig("demo_results/short_" + str(n) + ".png", dpi=200)
# The connections between keypoints are computed via the mid-range offsets.
# We can visuzalize them as well; for example right shoulder -> right hip
visualize_mid_offsets(offsets=sample_output[2],
heatmaps=H,
from_kp='Rshoulder',
to_kp='Rhip',
img=img,
every=8,
save_path=save_path)
# And we can see the reverse connection (Rhip -> Rshjoulder) as well
# visualize_mid_offsets(offsets= sample_output[2], heatmaps=H, to_kp='Rshoulder', from_kp='Rhip', img=img, every=8,save_path=save_path)
# We can use the heatmaps to compute the skeletons
pred_kp = get_keypoints(H)
print(pred_kp)
pred_skels = group_skeletons(keypoints=pred_kp, mid_offsets=sample_output[2])
pred_skels = [skel for skel in pred_skels if (skel[:, 2] > 0).sum() > 4]
print('Number of detected skeletons: {}'.format(len(pred_skels)))
plot_poses(img, pred_skels, save_path=save_path)
# we can use the predicted skeletons along with the long-range offsets and binary segmentation mask to compute the instance masks.
plt.imsave(
save_path + 'segmentation_mask.jpg',
apply_mask(img, sample_output[4][:, :, 0] > 0.5, color=[255, 255, 200]))
visualize_long_offsets(offsets=sample_output[3],
keypoint_id='Rshoulder',
seg_mask=sample_output[4],
def process_image(filepath, save_dir):
original = cv2.imread(filepath)
scale_outputs = []
for i in range(len(multiscale)):
scale = multiscale[i]
scale_img = get_multi_scale_img(img=original, scale=scale)
if i == 0:
img = scale_img[:, :, [2, 1, 0]]
plt.imsave(os.path.join(save_dir, 'input_image.jpg'), img)
imgs_batch = np.zeros(
(batch_size, int(scale * height), int(scale * width), 3))
imgs_batch[0] = scale_img
# make prediction
one_scale_output = sess.run(outputs[i],
feed_dict={tf_img[i]: imgs_batch})
scale_outputs.append([o[0] for o in one_scale_output])
sample_output = scale_outputs[0]
for i in range(1, len(multiscale)):
for j in range(len(sample_output)):
sample_output[j] += scale_outputs[i][j]
for j in range(len(sample_output)):
sample_output[j] /= len(multiscale)
# visualization
print('Visualization image has been saved into ', save_dir)
# Here is the output map for right shoulder
#Rshoulder_map = sample_output[0][:,:,config.KEYPOINTS.index('Rshoulder')]
#plt.imsave(save_path+'kp_map.jpg',overlay(img, Rshoulder_map, alpha=0.7))
# Gaussian filtering helps when there are multiple local maxima for the same keypoint.
H = compute_heatmaps(kp_maps=sample_output[0],
short_offsets=sample_output[1])
for i in range(17):
H[:, :, i] = gaussian_filter(H[:, :, i], sigma=2)
#plt.imsave(save_path+'heatmaps.jpg',H[:,:,config.KEYPOINTS.index('Rshoulder')]*10)
# The heatmaps are computed using the short offsets predicted by the network
# Here are the right shoulder offsets
#visualize_short_offsets(offsets=sample_output[1], heatmaps=H, keypoint_id='Rshoulder', img=img, every=8,save_path=save_path)
# The connections between keypoints are computed via the mid-range offsets.
# We can visuzalize them as well; for example right shoulder -> right hip
#visualize_mid_offsets(offsets= sample_output[2], heatmaps=H, from_kp='Rshoulder', to_kp='Rhip', img=img, every=8,save_path=save_path)
# And we can see the reverse connection (Rhip -> Rshjoulder) as well
# visualize_mid_offsets(offsets= sample_output[2], heatmaps=H, to_kp='Rshoulder', from_kp='Rhip', img=img, every=8,save_path=save_path)
# We can use the heatmaps to compute the skeletons
pred_kp = get_keypoints(H)
pred_skels = group_skeletons(keypoints=pred_kp,
mid_offsets=sample_output[2])
pred_skels = [skel for skel in pred_skels if (skel[:, 2] > 0).sum() > 4]
print('Number of detected skeletons: {}'.format(len(pred_skels)))
plot_poses(img, pred_skels, save_path=save_dir)
# we can use the predicted skeletons along with the long-range offsets and binary segmentation mask to compute the instance masks.
applied_mask = apply_mask(img,
sample_output[4][:, :, 0] > 0.5,
color=[255, 0, 0])
plt.imsave(os.path.join(save_dir, 'segmentation_mask.jpg'), applied_mask)
mask = sample_output[4][:, :, 0] > 0.5
temp = np.zeros(img.shape, dtype=np.uint8)
temp.fill(255)
for c in range(3):
temp[:, :, c] = np.where(mask == 1, 255, 0)
plt.imsave(os.path.join(save_dir, 'mask_new_scale.jpg'), temp)
temp = crop_to_original_size(original, temp)
plt.imsave(os.path.join(save_dir, 'MASK.jpg'), temp)
#visualize_long_offsets(offsets=sample_output[3], keypoint_id='Rshoulder', seg_mask=sample_output[4], img=img, every=8,save_path=save_path)
if len(pred_kp) > 0:
instance_masks = get_instance_masks(pred_skels, sample_output[-1][:, :,
0],
sample_output[-2])
plot_instance_masks(instance_masks, img, save_path=save_dir)
def visualize_mid_offsets(offsets,
from_kp,
to_kp,
centers=None,
heatmaps=None,
radius=config.KP_RADIUS,
img=None,
every=1,
save_path='./'):
if centers is None and heatmaps is None:
raise ValueError(
'either keypoint locations or heatmaps must be provided')
if isinstance(from_kp, str):
if not from_kp in config.KEYPOINTS:
raise ValueError('{} not a valid keypoint name'.format(from_kp))
else:
from_kp = config.KEYPOINTS.index(from_kp)
if isinstance(to_kp, str):
if not to_kp in config.KEYPOINTS:
raise ValueError('{} not a valid keypoint name'.format(to_kp))
else:
to_kp = config.KEYPOINTS.index(to_kp)
edge_list = config.EDGES + [edge[::-1] for edge in config.EDGES]
edge_id = edge_list.index((from_kp, to_kp))
if centers is None:
kp = get_keypoints(heatmaps)
kp = [k for k in kp if k['id'] == from_kp]
centers = [k['xy'].tolist() for k in kp]
kp_offsets = offsets[:, :, 2 * edge_id:2 * edge_id + 2]
# dists = np.zeros(offsets.shape[:2]+(len(centers),))
masks = np.zeros(offsets.shape[:2] + (len(centers), ), dtype='bool')
idx = np.rollaxis(np.indices(offsets.shape[1::-1]), 0, 3).transpose(
(1, 0, 2))
for j, c in enumerate(centers):
dists = np.sqrt(np.square(idx - c).sum(axis=-1))
dists_x = np.abs(idx[:, :, 0] - c[0])
dists_y = np.abs(idx[:, :, 1] - c[1])
masks[:, :, j] = (dists <= radius)
if every > 1:
d_mask = np.logical_and(
np.mod(dists_x.astype('int32'), every) == 0,
np.mod(dists_y.astype('int32'), every) == 0)
masks[:, :, j] = np.logical_and(masks[:, :, j], d_mask)
mask = masks.sum(axis=-1) > 0
# dists = dists.min(axis=-1)
# mask = dists <= radius
I, J = np.nonzero(mask)
if img is not None:
plt.imshow(img)
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['figure.dpi'] = 200
plt.quiver(J,
I,
kp_offsets[I, J, 0],
kp_offsets[I, J, 1],
color='r',
angles='xy',
scale_units='xy',
scale=1)
plt.savefig(save_path + 'middle_offsets.jpg', bbox_inches='tight')
