def accuracy(output, target, hm_type='gaussian', thr=0.5):
'''
Calculate accuracy according to PCK,
but uses ground truth heatmap rather than x,y locations
First value to be returned is average accuracy across 'idxs',
followed by individual accuracies
'''
idx = list(range(output.shape[1]))
norm = 1.0
if hm_type == 'gaussian':
pred, _ = get_max_preds(output)
target, _ = get_max_preds(target)
h = output.shape[2]
w = output.shape[3]
norm = np.ones((pred.shape[0], 2)) * np.array([h, w]) / 10
dists = calc_dists(pred, target, norm)
acc = np.zeros((len(idx) + 1))
avg_acc = 0
cnt = 0
for i in range(len(idx)):
acc[i + 1] = dist_acc(dists[idx[i]])
if acc[i + 1] >= 0:
avg_acc = avg_acc + acc[i + 1]
cnt += 1
avg_acc = avg_acc / cnt if cnt != 0 else 0
if cnt != 0:
acc[0] = avg_acc
return acc, avg_acc, cnt, pred
def detect_pose_secway(self, im, bbox):
cords = bbox[0:4]
H, W, C = im.shape
xmin, ymin, xmax, ymax = [int(i) for i in cords]
xmax = min(xmax, W)
ymax = min(ymax, H)
xmin = max(xmin, 0)
ymin = max(ymin, 0)
w, h = xmax - xmin, ymax - ymin
if w < 0 or h < 0:
return np.zeros([17, 2]), np.zeros([17]), np.zeros([17, 64, 64])
im_in = im[ymin:ymax, xmin:xmax, :]
scale_factor = [self.image_width / w, self.image_height / h]
im_resize = cv2.resize(im_in,
(int(self.image_width), int(self.image_height)),
interpolation=cv2.INTER_LINEAR)
with torch.no_grad():
input = self.transform(im_resize)
input = input.unsqueeze(0)
with torch.cuda.device(self.gpu_id):
adj = self.reset_adj_no1()
adj = adj.unsqueeze(0)
adj = adj.type(torch.cuda.FloatTensor)
input = input.cuda()
output = self.model(input, adj)
if cfg.TEST.FLIP_TEST:
input_flipped = np.flip(input.cpu().numpy(), 3).copy()
input_flipped = torch.from_numpy(input_flipped).cuda()
output_flipped = self.model(input_flipped, adj)
output_flipped = flip_back(output_flipped.cpu().numpy(),
self.flip_pairs)
output_flipped = torch.from_numpy(
output_flipped.copy()).cuda()
if cfg.TEST.SHIFT_HEATMAP:
output_flipped[:, :, :,
1:] = output_flipped.clone()[:, :, :,
0:-1]
output = (output + output_flipped) * 0.5
preds, maxvals = get_max_preds(output.clone().cpu().numpy())
preds, maxvals = preds.squeeze() * 4, maxvals.squeeze()
heatmaps = output.cpu().numpy().squeeze()
#For posetrack dataset, the 3th and 4th channel is None
preds = np.delete(preds, [3, 4], axis=0)
preds[:, 0] = preds[:, 0] / scale_factor[0] + bbox[0]
preds[:, 1] = preds[:, 1] / scale_factor[1] + bbox[1]
maxvals = np.delete(maxvals, [3, 4], axis=0)
heatmaps = np.delete(heatmaps, [3, 4], axis=0)
return preds, maxvals, heatmaps
def get_final_preds(self, batch_heatmaps, center, scale): coords, maxvals = get_max_preds(batch_heatmaps) heatmap_height = batch_heatmaps.shape[2] heatmap_width = batch_heatmaps.shape[3] # post-processing if cfg.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 preds = coords.copy() for i in range(coords.shape[0]): preds[i] = transform_preds(coords[i], center[i], scale[i], [heatmap_width, heatmap_height]) return preds, maxvals
def save_batch_heatmaps(batch_image, batch_heatmaps, file_name,
normalize=True):
'''
batch_image: [batch_size, channel, height, width]
batch_heatmaps: ['batch_size, num_joints, height, width]
file_name: saved file name
'''
if normalize:
batch_image = batch_image.clone()
min = float(batch_image.min())
max = float(batch_image.max())
batch_image.add_(-min).div_(max - min + 1e-5)
batch_size = batch_heatmaps.size(0)
num_joints = batch_heatmaps.size(1)
heatmap_height = batch_heatmaps.size(2)
heatmap_width = batch_heatmaps.size(3)
grid_image = np.zeros((batch_size*heatmap_height,
(num_joints+1)*heatmap_width,
3),
dtype=np.uint8)
preds, maxvals = get_max_preds(batch_heatmaps.detach().cpu().numpy())
for i in range(batch_size):
image = batch_image[i].mul(255)\
.clamp(0, 255)\
.byte()\
.permute(1, 2, 0)\
.cpu().numpy()
heatmaps = batch_heatmaps[i].mul(255)\
.clamp(0, 255)\
.byte()\
.cpu().numpy()
resized_image = cv2.resize(image,
(int(heatmap_width), int(heatmap_height)))
height_begin = heatmap_height * i
height_end = heatmap_height * (i + 1)
for j in range(num_joints):
cv2.circle(resized_image,
(int(preds[i][j][0]), int(preds[i][j][1])),
1, [0, 0, 255], 1)
heatmap = heatmaps[j, :, :]
colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
masked_image = colored_heatmap*0.7 + resized_image*0.3
cv2.circle(masked_image,
(int(preds[i][j][0]), int(preds[i][j][1])),
1, [0, 0, 255], 1)
width_begin = heatmap_width * (j+1)
width_end = heatmap_width * (j+2)
grid_image[height_begin:height_end, width_begin:width_end, :] = \
masked_image
# grid_image[height_begin:height_end, width_begin:width_end, :] = \
# colored_heatmap*0.7 + resized_image*0.3
grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image
cv2.imwrite(file_name, grid_image)