File size: 1,031 Bytes
8390f90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
r""" Evaluates CHMNet with PCK """
import torch
class Evaluator:
r""" Computes evaluation metrics of PCK """
@classmethod
def initialize(cls, alpha):
cls.alpha = torch.tensor(alpha).unsqueeze(1)
@classmethod
def evaluate(cls, prd_kps, batch):
r""" Compute percentage of correct key-points (PCK) with multiple alpha {0.05, 0.1, 0.15 }"""
pcks = []
for idx, (pk, tk) in enumerate(zip(prd_kps, batch['trg_kps'])):
pckthres = batch['pckthres'][idx]
npt = batch['n_pts'][idx]
prd_kps = pk[:, :npt]
trg_kps = tk[:, :npt]
l2dist = (prd_kps - trg_kps).pow(2).sum(dim=0).pow(0.5).unsqueeze(0).repeat(len(cls.alpha), 1)
thres = pckthres.expand_as(l2dist).float() * cls.alpha
pck = torch.le(l2dist, thres).sum(dim=1) / float(npt)
if len(pck) == 1: pck = pck[0]
pcks.append(pck)
eval_result = {'pck': pcks}
return eval_result
|