|
|
|
|
|
|
|
|
|
|
|
|
| from __future__ import absolute_import
|
| from __future__ import division
|
| from __future__ import print_function
|
|
|
| import math
|
|
|
| import numpy as np
|
|
|
|
|
| def get_max_preds(batch_heatmaps):
|
| '''
|
| get predictions from score maps
|
| heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
|
| '''
|
| assert isinstance(batch_heatmaps, np.ndarray), \
|
| 'batch_heatmaps should be numpy.ndarray'
|
| assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'
|
|
|
| batch_size = batch_heatmaps.shape[0]
|
| num_joints = batch_heatmaps.shape[1]
|
| width = batch_heatmaps.shape[3]
|
| heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))
|
| idx = np.argmax(heatmaps_reshaped, 2)
|
| maxvals = np.amax(heatmaps_reshaped, 2)
|
|
|
| maxvals = maxvals.reshape((batch_size, num_joints, 1))
|
| idx = idx.reshape((batch_size, num_joints, 1))
|
|
|
| preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
|
|
|
| preds[:, :, 0] = (preds[:, :, 0]) % width
|
| preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)
|
|
|
| pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
|
| pred_mask = pred_mask.astype(np.float32)
|
|
|
| preds *= pred_mask
|
| return preds, maxvals
|
|
|
|
|
| def get_final_preds(config, output, center, scale, coord_heatmaps=None):
|
| heatmap_height = config.MODEL.HEATMAP_SIZE[1]
|
| heatmap_width = config.MODEL.HEATMAP_SIZE[0]
|
|
|
| batch_heatmaps = output
|
| coords, maxvals = get_max_preds(batch_heatmaps)
|
|
|
| if config.TEST.POST_PROCESS:
|
| for n in range(coords.shape[0]):
|
| for p in range(coords.shape[1]):
|
| hm = batch_heatmaps[n][p]
|
| px = int(math.floor(coords[n][p][0] + 0.5))
|
| py = int(math.floor(coords[n][p][1] + 0.5))
|
| if 1 < px < heatmap_width-1 and 1 < py < heatmap_height-1:
|
| diff = np.array(
|
| [
|
| hm[py][px+1] - hm[py][px-1],
|
| hm[py+1][px]-hm[py-1][px]
|
| ]
|
| )
|
| coords[n][p] += np.sign(diff) * .25
|
|
|
| return coords, maxvals
|
|
|
