94 lines
3.5 KiB
Python
94 lines
3.5 KiB
Python
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: {}'.format(batch_heatmaps.shape)
|
|
|
|
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
|
|
|
|
COCO17_IN_BODY25 = [0,16,15,18,17,5,2,6,3,7,4,12,9,13,10,14,11]
|
|
pairs = [[1, 8], [1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], [8, 9], [9, 10], [10, 11], [8, 12], [12, 13], [13, 14], [1, 0], [0,15], [15,17], [0,16], [16,18], [14,19], [19,20], [14,21], [11,22], [22,23], [11,24]]
|
|
def coco17tobody25(points2d):
|
|
kpts = np.zeros((points2d.shape[0], 25, 3))
|
|
kpts[:, COCO17_IN_BODY25, :2] = points2d[:, :, :2]
|
|
kpts[:, COCO17_IN_BODY25, 2:3] = points2d[:, :, 2:3]
|
|
kpts[:, 8, :2] = kpts[:, [9, 12], :2].mean(axis=1)
|
|
kpts[:, 8, 2] = kpts[:, [9, 12], 2].min(axis=1)
|
|
kpts[:, 1, :2] = kpts[:, [2, 5], :2].mean(axis=1)
|
|
kpts[:, 1, 2] = kpts[:, [2, 5], 2].min(axis=1)
|
|
# 需要交换一下
|
|
# kpts = kpts[:, :, [1,0,2]]
|
|
return kpts
|
|
|
|
def coco23tobody25(points2d):
|
|
kpts = coco17tobody25(points2d[:, :17])
|
|
kpts[:, [19, 20, 21, 22, 23, 24]] = points2d[:, [17, 18, 19, 20, 21, 22]]
|
|
return kpts
|
|
|
|
class BaseKeypoints():
|
|
@staticmethod
|
|
def get_max_preds(batch_heatmaps):
|
|
coords, maxvals = get_max_preds(batch_heatmaps)
|
|
|
|
heatmap_height = batch_heatmaps.shape[2]
|
|
heatmap_width = batch_heatmaps.shape[3]
|
|
|
|
# post-processing
|
|
if True:
|
|
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
|
|
coords = coords.astype(np.float32) * 4
|
|
pred = np.dstack((coords, maxvals))
|
|
return pred
|
|
|
|
@staticmethod
|
|
def batch_affine_transform(points, trans):
|
|
# points: (Bn, J, 2), trans: (Bn, 2, 3)
|
|
points = np.dstack((points[..., :2], np.ones((*points.shape[:-1], 1))))
|
|
out = np.matmul(points, trans.swapaxes(-1, -2))
|
|
return out
|
|
|
|
@staticmethod
|
|
def coco17tobody25(points2d):
|
|
return coco17tobody25(points2d)
|
|
|
|
@staticmethod
|
|
def coco23tobody25(points2d):
|
|
return coco23tobody25(points2d)
|
|
|