Spaces:
Running
Running
| from opt import opt | |
| try: | |
| from utils.img import transformBoxInvert, transformBoxInvert_batch, findPeak, processPeaks | |
| except ImportError: | |
| from SPPE.src.utils.img import transformBoxInvert, transformBoxInvert_batch, findPeak, processPeaks | |
| import torch | |
| class DataLogger(object): | |
| def __init__(self): | |
| self.clear() | |
| def clear(self): | |
| self.value = 0 | |
| self.sum = 0 | |
| self.cnt = 0 | |
| self.avg = 0 | |
| def update(self, value, n=1): | |
| self.value = value | |
| self.sum += value * n | |
| self.cnt += n | |
| self._cal_avg() | |
| def _cal_avg(self): | |
| self.avg = self.sum / self.cnt | |
| def accuracy(output, label, dataset): | |
| if type(output) == list: | |
| return accuracy(output[opt.nStack - 1], label[opt.nStack - 1], dataset) | |
| else: | |
| return heatmapAccuracy(output.cpu().data, label.cpu().data, dataset.accIdxs) | |
| def heatmapAccuracy(output, label, idxs): | |
| preds = getPreds(output) | |
| gt = getPreds(label) | |
| norm = torch.ones(preds.size(0)) * opt.outputResH / 10 | |
| dists = calc_dists(preds, gt, norm) | |
| #print(dists) | |
| acc = torch.zeros(len(idxs) + 1) | |
| avg_acc = 0 | |
| cnt = 0 | |
| for i in range(len(idxs)): | |
| acc[i + 1] = dist_acc(dists[idxs[i] - 1]) | |
| if acc[i + 1] >= 0: | |
| avg_acc = avg_acc + acc[i + 1] | |
| cnt += 1 | |
| if cnt != 0: | |
| acc[0] = avg_acc / cnt | |
| return acc | |
| def getPreds(hm): | |
| ''' get predictions from score maps in torch Tensor | |
| return type: torch.LongTensor | |
| ''' | |
| assert hm.dim() == 4, 'Score maps should be 4-dim' | |
| maxval, idx = torch.max(hm.view(hm.size(0), hm.size(1), -1), 2) | |
| maxval = maxval.view(hm.size(0), hm.size(1), 1) | |
| idx = idx.view(hm.size(0), hm.size(1), 1) + 1 | |
| preds = idx.repeat(1, 1, 2).float() | |
| preds[:, :, 0] = (preds[:, :, 0] - 1) % hm.size(3) | |
| preds[:, :, 1] = torch.floor((preds[:, :, 1] - 1) / hm.size(3)) | |
| # pred_mask = maxval.gt(0).repeat(1, 1, 2).float() | |
| # preds *= pred_mask | |
| return preds | |
| def calc_dists(preds, target, normalize): | |
| preds = preds.float().clone() | |
| target = target.float().clone() | |
| dists = torch.zeros(preds.size(1), preds.size(0)) | |
| for n in range(preds.size(0)): | |
| for c in range(preds.size(1)): | |
| if target[n, c, 0] > 0 and target[n, c, 1] > 0: | |
| dists[c, n] = torch.dist( | |
| preds[n, c, :], target[n, c, :]) / normalize[n] | |
| else: | |
| dists[c, n] = -1 | |
| return dists | |
| def dist_acc(dists, thr=0.5): | |
| ''' Return percentage below threshold while ignoring values with a -1 ''' | |
| if dists.ne(-1).sum() > 0: | |
| return dists.le(thr).eq(dists.ne(-1)).float().sum() * 1.0 / dists.ne(-1).float().sum() | |
| else: | |
| return - 1 | |
| def postprocess(output): | |
| p = getPreds(output) | |
| for i in range(p.size(0)): | |
| for j in range(p.size(1)): | |
| hm = output[i][j] | |
| pX, pY = int(round(p[i][j][0])), int(round(p[i][j][1])) | |
| if 0 < pX < opt.outputResW - 1 and 0 < pY < opt.outputResH - 1: | |
| diff = torch.Tensor((hm[pY][pX + 1] - hm[pY][pX - 1], hm[pY + 1][pX] - hm[pY - 1][pX])) | |
| p[i][j] += diff.sign() * 0.25 | |
| p -= 0.5 | |
| return p | |
| def getPrediction(hms, pt1, pt2, inpH, inpW, resH, resW): | |
| ''' | |
| Get keypoint location from heatmaps | |
| ''' | |
| assert hms.dim() == 4, 'Score maps should be 4-dim' | |
| maxval, idx = torch.max(hms.view(hms.size(0), hms.size(1), -1), 2) | |
| maxval = maxval.view(hms.size(0), hms.size(1), 1) | |
| idx = idx.view(hms.size(0), hms.size(1), 1) + 1 | |
| preds = idx.repeat(1, 1, 2).float() | |
| preds[:, :, 0] = (preds[:, :, 0] - 1) % hms.size(3) | |
| preds[:, :, 1] = torch.floor((preds[:, :, 1] - 1) / hms.size(3)) | |
| pred_mask = maxval.gt(0).repeat(1, 1, 2).float() | |
| preds *= pred_mask | |
| # Very simple post-processing step to improve performance at tight PCK thresholds | |
| for i in range(preds.size(0)): | |
| for j in range(preds.size(1)): | |
| hm = hms[i][j] | |
| pX, pY = int(round(float(preds[i][j][0]))), int(round(float(preds[i][j][1]))) | |
| if 0 < pX < opt.outputResW - 1 and 0 < pY < opt.outputResH - 1: | |
| diff = torch.Tensor( | |
| (hm[pY][pX + 1] - hm[pY][pX - 1], hm[pY + 1][pX] - hm[pY - 1][pX])) | |
| preds[i][j] += diff.sign() * 0.25 | |
| preds += 0.2 | |
| preds_tf = torch.zeros(preds.size()) | |
| preds_tf = transformBoxInvert_batch(preds, pt1, pt2, inpH, inpW, resH, resW) | |
| return preds, preds_tf, maxval | |
| def getMultiPeakPrediction(hms, pt1, pt2, inpH, inpW, resH, resW): | |
| assert hms.dim() == 4, 'Score maps should be 4-dim' | |
| preds_img = {} | |
| hms = hms.numpy() | |
| for n in range(hms.shape[0]): # Number of samples | |
| preds_img[n] = {} # Result of sample: n | |
| for k in range(hms.shape[1]): # Number of keypoints | |
| preds_img[n][k] = [] # Result of keypoint: k | |
| hm = hms[n][k] | |
| candidate_points = findPeak(hm) | |
| res_pt = processPeaks(candidate_points, hm, | |
| pt1[n], pt2[n], inpH, inpW, resH, resW) | |
| preds_img[n][k] = res_pt | |
| return preds_img | |
| def getPrediction_batch(hms, pt1, pt2, inpH, inpW, resH, resW): | |
| ''' | |
| Get keypoint location from heatmaps | |
| pt1, pt2: [n, 2] | |
| OUTPUT: | |
| preds: [n, 17, 2] | |
| ''' | |
| assert hms.dim() == 4, 'Score maps should be 4-dim' | |
| flat_hms = hms.view(hms.size(0), hms.size(1), -1) | |
| maxval, idx = torch.max(flat_hms, 2) | |
| maxval = maxval.view(hms.size(0), hms.size(1), 1) | |
| idx = idx.view(hms.size(0), hms.size(1), 1) + 1 | |
| preds = idx.repeat(1, 1, 2).float() | |
| preds[:, :, 0] = (preds[:, :, 0] - 1) % hms.size(3) | |
| preds[:, :, 1] = torch.floor((preds[:, :, 1] - 1) / hms.size(3)) | |
| pred_mask = maxval.gt(0).repeat(1, 1, 2).float() | |
| preds *= pred_mask | |
| # Very simple post-processing step to improve performance at tight PCK thresholds | |
| idx_up = (idx - hms.size(3)).clamp(0, flat_hms.size(2) - 1) | |
| idx_down = (idx + hms.size(3)).clamp(0, flat_hms.size(2) - 1) | |
| idx_left = (idx - 1).clamp(0, flat_hms.size(2) - 1) | |
| idx_right = (idx + 1).clamp(0, flat_hms.size(2) - 1) | |
| maxval_up = flat_hms.gather(2, idx_up) | |
| maxval_down = flat_hms.gather(2, idx_down) | |
| maxval_left = flat_hms.gather(2, idx_left) | |
| maxval_right = flat_hms.gather(2, idx_right) | |
| diff1 = (maxval_right - maxval_left).sign() * 0.25 | |
| diff2 = (maxval_down - maxval_up).sign() * 0.25 | |
| diff1[idx_up <= hms.size(3)] = 0 | |
| diff1[idx_down / hms.size(3) >= (hms.size(3) - 1)] = 0 | |
| diff2[(idx_left % hms.size(3)) == 0] = 0 | |
| diff2[(idx_left % hms.size(3)) == (hms.size(3) - 1)] = 0 | |
| preds[:, :, 0] += diff1.squeeze(-1) | |
| preds[:, :, 1] += diff2.squeeze(-1) | |
| preds_tf = torch.zeros(preds.size()) | |
| preds_tf = transformBoxInvert_batch(preds, pt1, pt2, inpH, inpW, resH, resW) | |
| return preds, preds_tf, maxval | |