def forward(self, pred, target):
ova_len = len(pred) * self.n_kps
if self.pb_type == 'regression':
shouders_len = (
(target[..., self.sr:self.sr + 1] -
target[..., self.sl:self.sl + 1])**2 +
(target[..., self.sr + self.n_kps:self.sr + self.n_kps + 1] -
target[..., self.sl + self.n_kps:self.sl + self.n_kps + 1])**
2)**0.5
err = torch.abs(pred - target)
err = (err[..., :self.n_kps]**2 + err[..., self.n_kps]**2)**0.5
err = torch.sum(err < shouders_len * self.thresh)
elif self.pb_type == 'detection':
pred = heatmap2coor(pred, self.n_kps, self.img_size, self.stride)
target = heatmap2coor(target, self.n_kps, self.img_size,
self.stride)
shouders_len = ((target[:, self.sr:self.sr + 1, 0] -
target[:, self.sl:self.sl + 1, 0])**2 +
(target[:, self.sr:self.sr + 1, 1] -
target[:, self.sl:self.sl + 1, 1])**2)**0.5
err = torch.abs(pred - target)
err = (err[..., 0]**2 + err[..., 1]**2)**0.5
err = torch.sum(err < shouders_len * self.thresh)
else:
return None
return err / ova_len
def forward(self, pred, target):
if self.pb_type == 'regression':
ova_loss = torch.mean(torch.sum(torch.abs(pred-target), dim=-1)/(2*self.n_kps))
elif self.pb_type == 'detection':
pred = heatmap2coor(pred, self.n_kps, self.img_size, self.stride)
target = heatmap2coor(target, self.n_kps, self.img_size, self.stride)
ova_loss = torch.mean(torch.sum(torch.abs(pred-target), dim=(-1,-2))/(2*self.n_kps))
else:
return None
return ova_loss
def predict(self, img):
self.model.eval()
with torch.no_grad() as tng:
preds = self.model(img)
preds = preds.cpu()
if self.pb_type == 'regression':
preds = preds.numpy()
preds = np.stack([preds[:, ::2], preds[:, 1:][:, ::2]],
axis=-1)
elif self.pb_type == 'detection' or self.pb_type == 'define':
preds = heatmap2coor(preds, self.n_kps, self.img_size,
self.stride).numpy()
return preds