def calc_stress_err()

in dataset/eval_zoo.py [0:0]

• 38 lines of code
• 8 McCabe index (conditional complexity)

def calc_stress_err(gt, pred, scale=1., mask=None, get_best_scale=False):

    assert pred.shape[1] == 3
    assert gt.shape[1] == 3
    assert pred.shape[0] == gt.shape[0]
    assert pred.shape[2] == gt.shape[2]

    if get_best_scale:
        argmin_scale = compute_best_scale(pred, gt, v=mask)
        pred = pred.copy() * argmin_scale[:, None, None]

    errs = []
    nkp = gt.shape[2]
    if mask is not None:
        tridx_cache = [np.triu_indices(k, k=1)
                       for k in range(nkp+1)]
        assert mask.shape[1] == pred.shape[2]
        assert mask.shape[0] == pred.shape[0]
    else:
        tridx = np.triu_indices(nkp, k=1)
        assert len(tridx[0]) == (nkp*(nkp-1))/2

    print('stress eval:')
    with tqdm(total=len(gt)) as tq:
        for ii, (g_, p_) in enumerate(zip(gt, pred)):
            if mask is not None:
                mask_ = mask[ii]
            else:
                mask_ = None
            edm_g = calc_edm(g_, squared=False, mask=mask_)
            edm_p = calc_edm(p_, squared=False, mask=mask_)
            stress = np.abs(edm_g - edm_p)
            if mask_ is not None:
                nkp_now = edm_g.shape[0]
                assert mask_.sum() == nkp_now
                tridx = tridx_cache[nkp_now]
            mstress = stress[tridx[0], tridx[1]]
            mstress = mstress.mean()
            # if True:
            # 	triu_mask_ = np.triu(np.ones(nkp),k=1)
            # 	mstress_ = (stress * triu_mask_).sum() / triu_mask_.sum()
            # 	assert np.abs(mstress - mstress_) <= 1e-3

            errs.append(mstress)
            tq.update(1)

    errs = np.array(errs) * scale

    return errs