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')