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 getPrediction_batch(hms, pt1, pt2, inpH, inpW, resH, resW):
'''
Get keypoint location from heatmaps
pt1, pt2: [n, 2]
OUTPUT:
preds: [n, 50, 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
def adjustPrediction(hms,pt1,pt2,inpH,inpW,resH,resW):
assert hms.dim() == 3, 'predict maps should be 3-dim'
preds = hms[:,...,:2]
maxval = hms[:,...,2:]
# Very simple post-processing step to improve performance at tight PCK thresholds
preds_tf = torch.zeros(preds.size())
preds_tf = transformBoxInvert_batch(preds, pt1, pt2, inpH, inpW, resH, resW)
return preds, preds_tf, maxval
def getModifiyPosition(preds,hms, pt1, pt2, inpH, inpW, resH, resW):
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