import cv2 import math import numpy as np from preprocess_ops import get_affine_transform class HRNetPostProcess(object): def __init__(self, use_dark=True): self.use_dark = use_dark def flip_back(self, output_flipped, matched_parts): assert output_flipped.ndim == 4,\ 'output_flipped should be [batch_size, num_joints, height, width]' output_flipped = output_flipped[:, :, :, ::-1] for pair in matched_parts: tmp = output_flipped[:, pair[0], :, :].copy() output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :] output_flipped[:, pair[1], :, :] = tmp return output_flipped def get_max_preds(self, heatmaps): """get predictions from score maps Args: heatmaps: numpy.ndarray([batch_size, num_joints, height, width]) Returns: preds: numpy.ndarray([batch_size, num_joints, 2]), keypoints coords maxvals: numpy.ndarray([batch_size, num_joints, 2]), the maximum confidence of the keypoints """ assert isinstance(heatmaps, np.ndarray), 'heatmaps should be numpy.ndarray' assert heatmaps.ndim == 4, 'batch_images should be 4-ndim' batch_size = heatmaps.shape[0] num_joints = heatmaps.shape[1] width = heatmaps.shape[3] heatmaps_reshaped = 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 gaussian_blur(self, heatmap, kernel): border = (kernel - 1) // 2 batch_size = heatmap.shape[0] num_joints = heatmap.shape[1] height = heatmap.shape[2] width = heatmap.shape[3] for i in range(batch_size): for j in range(num_joints): origin_max = np.max(heatmap[i, j]) dr = np.zeros((height + 2 * border, width + 2 * border)) dr[border:-border, border:-border] = heatmap[i, j].copy() dr = cv2.GaussianBlur(dr, (kernel, kernel), 0) heatmap[i, j] = dr[border:-border, border:-border].copy() heatmap[i, j] *= origin_max / np.max(heatmap[i, j]) return heatmap def dark_parse(self, hm, coord): heatmap_height = hm.shape[0] heatmap_width = hm.shape[1] px = int(coord[0]) py = int(coord[1]) if 1 < px < heatmap_width - 2 and 1 < py < heatmap_height - 2: dx = 0.5 * (hm[py][px + 1] - hm[py][px - 1]) dy = 0.5 * (hm[py + 1][px] - hm[py - 1][px]) dxx = 0.25 * (hm[py][px + 2] - 2 * hm[py][px] + hm[py][px - 2]) dxy = 0.25 * (hm[py+1][px+1] - hm[py-1][px+1] - hm[py+1][px-1] \ + hm[py-1][px-1]) dyy = 0.25 * ( hm[py + 2 * 1][px] - 2 * hm[py][px] + hm[py - 2 * 1][px]) derivative = np.matrix([[dx], [dy]]) hessian = np.matrix([[dxx, dxy], [dxy, dyy]]) if dxx * dyy - dxy**2 != 0: hessianinv = hessian.I offset = -hessianinv * derivative offset = np.squeeze(np.array(offset.T), axis=0) coord += offset return coord def dark_postprocess(self, hm, coords, kernelsize): """ refer to https://github.com/ilovepose/DarkPose/lib/core/inference.py """ hm = self.gaussian_blur(hm, kernelsize) hm = np.maximum(hm, 1e-10) hm = np.log(hm) for n in range(coords.shape[0]): for p in range(coords.shape[1]): coords[n, p] = self.dark_parse(hm[n][p], coords[n][p]) return coords def get_final_preds(self, heatmaps, center, scale, kernelsize=3): """the highest heatvalue location with a quarter offset in the direction from the highest response to the second highest response. Args: heatmaps (numpy.ndarray): The predicted heatmaps center (numpy.ndarray): The boxes center scale (numpy.ndarray): The scale factor Returns: preds: numpy.ndarray([batch_size, num_joints, 2]), keypoints coords maxvals: numpy.ndarray([batch_size, num_joints, 1]), the maximum confidence of the keypoints """ coords, maxvals = self.get_max_preds(heatmaps) heatmap_height = heatmaps.shape[2] heatmap_width = heatmaps.shape[3] if self.use_dark: coords = self.dark_postprocess(heatmaps, coords, kernelsize) else: for n in range(coords.shape[0]): for p in range(coords.shape[1]): hm = 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() # Transform back 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 __call__(self, output, center, scale): preds, maxvals = self.get_final_preds(output, center, scale) return np.concatenate( (preds, maxvals), axis=-1), np.mean( maxvals, axis=1) def transform_preds(coords, center, scale, output_size): target_coords = np.zeros(coords.shape) trans = get_affine_transform(center, scale * 200, 0, output_size, inv=1) for p in range(coords.shape[0]): target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans) return target_coords def affine_transform(pt, t): new_pt = np.array([pt[0], pt[1], 1.]).T new_pt = np.dot(t, new_pt) return new_pt[:2]