Ejemplo n.º 1
0
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
Ejemplo n.º 2
0
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
Ejemplo n.º 3
0
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
Ejemplo n.º 4
0
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