File size: 3,048 Bytes
bbde80b |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import torch
import torch.nn as nn
import ipdb
from lib.utils.utils_mesh import batch_rodrigues
from lib.model.loss import *
class MeshLoss(nn.Module):
def __init__(
self,
loss_type='MSE',
device='cuda',
):
super(MeshLoss, self).__init__()
self.device = device
self.loss_type = loss_type
if loss_type == 'MSE':
self.criterion_keypoints = nn.MSELoss(reduction='none').to(self.device)
self.criterion_regr = nn.MSELoss().to(self.device)
elif loss_type == 'L1':
self.criterion_keypoints = nn.L1Loss(reduction='none').to(self.device)
self.criterion_regr = nn.L1Loss().to(self.device)
def forward(
self,
smpl_output,
data_gt,
):
# to reduce time dimension
reduce = lambda x: x.reshape((x.shape[0] * x.shape[1],) + x.shape[2:])
data_3d_theta = reduce(data_gt['theta'])
preds = smpl_output[-1]
pred_theta = preds['theta']
theta_size = pred_theta.shape[:2]
pred_theta = reduce(pred_theta)
preds_local = preds['kp_3d'] - preds['kp_3d'][:, :, 0:1,:] # (N, T, 17, 3)
gt_local = data_gt['kp_3d'] - data_gt['kp_3d'][:, :, 0:1,:]
real_shape, pred_shape = data_3d_theta[:, 72:], pred_theta[:, 72:]
real_pose, pred_pose = data_3d_theta[:, :72], pred_theta[:, :72]
loss_dict = {}
loss_dict['loss_3d_pos'] = loss_mpjpe(preds_local, gt_local)
loss_dict['loss_3d_scale'] = n_mpjpe(preds_local, gt_local)
loss_dict['loss_3d_velocity'] = loss_velocity(preds_local, gt_local)
loss_dict['loss_lv'] = loss_limb_var(preds_local)
loss_dict['loss_lg'] = loss_limb_gt(preds_local, gt_local)
loss_dict['loss_a'] = loss_angle(preds_local, gt_local)
loss_dict['loss_av'] = loss_angle_velocity(preds_local, gt_local)
if pred_theta.shape[0] > 0:
loss_pose, loss_shape = self.smpl_losses(pred_pose, pred_shape, real_pose, real_shape)
loss_norm = torch.norm(pred_theta, dim=-1).mean()
loss_dict['loss_shape'] = loss_shape
loss_dict['loss_pose'] = loss_pose
loss_dict['loss_norm'] = loss_norm
return loss_dict
def smpl_losses(self, pred_rotmat, pred_betas, gt_pose, gt_betas):
pred_rotmat_valid = batch_rodrigues(pred_rotmat.reshape(-1,3)).reshape(-1, 24, 3, 3)
gt_rotmat_valid = batch_rodrigues(gt_pose.reshape(-1,3)).reshape(-1, 24, 3, 3)
pred_betas_valid = pred_betas
gt_betas_valid = gt_betas
if len(pred_rotmat_valid) > 0:
loss_regr_pose = self.criterion_regr(pred_rotmat_valid, gt_rotmat_valid)
loss_regr_betas = self.criterion_regr(pred_betas_valid, gt_betas_valid)
else:
loss_regr_pose = torch.FloatTensor(1).fill_(0.).to(self.device)
loss_regr_betas = torch.FloatTensor(1).fill_(0.).to(self.device)
return loss_regr_pose, loss_regr_betas
|