def get_batch_result(self, type):
preds2d = get_preds_from_heatmap(self.outputs['heatmap'][-1])
# preds2d = get_preds_from_heatmap(self.batch['heatmap'])
preds3d = torch.zeros((preds2d.size(0), 21, 3))
root_depth = self.batch['root_depth']
index_bone_length = self.batch['index_bone_length']
preds3d[:, :, :2] = preds2d.clone()
preds3d[:, :, 2] = self.outputs['depth'] * index_bone_length.unsqueeze(
1) + root_depth.unsqueeze(1)
# preds3d[:,:,2] = self.batch['relative_depth'] * index_bone_length.unsqueeze(1) + root_depth.unsqueeze(1)
preds3d[:, :, :2] *= preds3d[:, :, 2:]
for i in range(preds3d.size(0)):
preds3d[i, :, :] = torch.matmul(
preds3d[i, :, :], self.batch['matrix'][i].transpose(0, 1))
dis2d = torch.norm(self.batch['coor2d'][..., :2] - preds2d, dim=-1)
dis3d = torch.norm(self.batch['coor3d'] - preds3d, dim=-1)
self.batch_result = {
'coor2d': preds2d,
'coor3d': preds3d,
'dis3d': dis3d,
'dis2d': dis2d
}
return self.batch_result
def get_preds(self):
preds2d = get_preds_from_heatmap(self.outputs['color_hm'][-1])
# preds2d = get_preds_from_heatmap(self.batch['heatmap'])
preds3d = torch.zeros((preds2d.size(0), 21, 3))
root_depth = self.batch['root_depth']
index_bone_length = self.batch['index_bone_length']
preds3d[:, :, :2] = preds2d.clone()
preds3d[:, :, 2] = self.outputs['depth'] * index_bone_length.unsqueeze(
1) + root_depth.unsqueeze(1)
# preds3d[:,:,2] = self.batch['relative_depth'] * index_bone_length.unsqueeze(1) + root_depth.unsqueeze(1)
preds3d[:, :, :2] *= preds3d[:, :, 2:]
for i in range(preds3d.size(0)):
preds3d[i, :, :] = torch.matmul(
preds3d[i, :, :], self.batch['matrix'][i].transpose(0, 1))
self.preds = {'pose2d': preds2d, 'pose3d': preds3d}
return self.preds
def get_preds(self):
preds_2d = get_preds_from_heatmap(self.outputs['heatmap'][-1])
return {'pose2d':preds_2d}
def eval_result(self):
preds_2d = get_preds_from_heatmap(self.outputs['heatmap'][-1])
dis_2d = torch.norm(self.batch['project'][:, :, :2] - preds_2d, dim = -1).mean()
return {"dis2d": dis_2d}