Spaces:
Runtime error
Runtime error
| import torch | |
| import numpy as np | |
| import pickle as pkl | |
| import os | |
| import sys | |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src')) | |
| # from priors.pose_prior_35 import Prior | |
| # from priors.tiger_pose_prior.tiger_pose_prior import GaussianMixturePrior | |
| from priors.normalizing_flow_prior.normalizing_flow_prior import NormalizingFlowPrior | |
| from priors.shape_prior import ShapePrior | |
| from lifting_to_3d.utils.geometry_utils import rot6d_to_rotmat, batch_rot2aa | |
| from configs.SMAL_configs import UNITY_SMAL_SHAPE_PRIOR_DOGS | |
| class Loss(torch.nn.Module): | |
| def __init__(self, data_info, nf_version=None): | |
| super(Loss, self).__init__() | |
| self.criterion_regr = torch.nn.MSELoss() # takes the mean | |
| self.criterion_class = torch.nn.CrossEntropyLoss() | |
| self.data_info = data_info | |
| self.register_buffer('keypoint_weights', torch.tensor(data_info.keypoint_weights)[None, :]) | |
| self.l_anchor = None | |
| self.l_pos = None | |
| self.l_neg = None | |
| if nf_version is not None: | |
| self.normalizing_flow_pose_prior = NormalizingFlowPrior(nf_version=nf_version) | |
| self.shape_prior = ShapePrior(UNITY_SMAL_SHAPE_PRIOR_DOGS) | |
| self.criterion_triplet = torch.nn.TripletMarginLoss(margin=1) | |
| # load 3d data for the unity dogs (an optional shape prior for 11 breeds) | |
| with open(UNITY_SMAL_SHAPE_PRIOR_DOGS, 'rb') as f: | |
| data = pkl.load(f) | |
| dog_betas_unity = data['dogs_betas'] | |
| self.dog_betas_unity = {29: torch.tensor(dog_betas_unity[0, :]).float(), | |
| 91: torch.tensor(dog_betas_unity[1, :]).float(), | |
| 84: torch.tensor(0.5*dog_betas_unity[3, :] + 0.5*dog_betas_unity[14, :]).float(), | |
| 85: torch.tensor(dog_betas_unity[5, :]).float(), | |
| 28: torch.tensor(dog_betas_unity[6, :]).float(), | |
| 94: torch.tensor(dog_betas_unity[7, :]).float(), | |
| 92: torch.tensor(dog_betas_unity[8, :]).float(), | |
| 95: torch.tensor(dog_betas_unity[10, :]).float(), | |
| 20: torch.tensor(dog_betas_unity[11, :]).float(), | |
| 83: torch.tensor(dog_betas_unity[12, :]).float(), | |
| 99: torch.tensor(dog_betas_unity[16, :]).float()} | |
| def prepare_anchor_pos_neg(self, batch_size, device): | |
| l0 = np.arange(0, batch_size, 2) | |
| l_anchor = [] | |
| l_pos = [] | |
| l_neg = [] | |
| for ind in l0: | |
| xx = set(np.arange(0, batch_size)) | |
| xx.discard(ind) | |
| xx.discard(ind+1) | |
| for ind2 in xx: | |
| if ind2 % 2 == 0: | |
| l_anchor.append(ind) | |
| l_pos.append(ind + 1) | |
| else: | |
| l_anchor.append(ind + 1) | |
| l_pos.append(ind) | |
| l_neg.append(ind2) | |
| self.l_anchor = torch.Tensor(l_anchor).to(torch.int64).to(device) | |
| self.l_pos = torch.Tensor(l_pos).to(torch.int64).to(device) | |
| self.l_neg = torch.Tensor(l_neg).to(torch.int64).to(device) | |
| return | |
| def forward(self, output_reproj, target_dict, weight_dict=None): | |
| # output_reproj: ['vertices_smal', 'keyp_3d', 'keyp_2d', 'silh_image'] | |
| # target_dict: ['index', 'center', 'scale', 'pts', 'tpts', 'target_weight'] | |
| batch_size = output_reproj['keyp_2d'].shape[0] | |
| # loss on reprojected keypoints | |
| output_kp_resh = (output_reproj['keyp_2d']).reshape((-1, 2)) | |
| target_kp_resh = (target_dict['tpts'][:, :, :2] / 64. * (256. - 1)).reshape((-1, 2)) | |
| weights_resh = target_dict['tpts'][:, :, 2].reshape((-1)) | |
| keyp_w_resh = self.keypoint_weights.repeat((batch_size, 1)).reshape((-1)) | |
| loss_keyp = ((((output_kp_resh - target_kp_resh)[weights_resh>0]**2).sum(axis=1).sqrt()*weights_resh[weights_resh>0])*keyp_w_resh[weights_resh>0]).sum() / \ | |
| max((weights_resh[weights_resh>0]*keyp_w_resh[weights_resh>0]).sum(), 1e-5) | |
| # loss on reprojected silhouette | |
| assert output_reproj['silh'].shape == (target_dict['silh'][:, None, :, :]).shape | |
| silh_loss_type = 'default' | |
| if silh_loss_type == 'default': | |
| with torch.no_grad(): | |
| thr_silh = 20 | |
| diff = torch.norm(output_kp_resh - target_kp_resh, dim=1) | |
| diff_x = diff.reshape((batch_size, -1)) | |
| weights_resh_x = weights_resh.reshape((batch_size, -1)) | |
| unweighted_kp_mean_dist = (diff_x * weights_resh_x).sum(dim=1) / ((weights_resh_x).sum(dim=1)+1e-6) | |
| loss_silh_bs = ((output_reproj['silh'] - target_dict['silh'][:, None, :, :]) ** 2).sum(axis=3).sum(axis=2).sum(axis=1) / (output_reproj['silh'].shape[2]*output_reproj['silh'].shape[3]) | |
| loss_silh = loss_silh_bs[unweighted_kp_mean_dist<thr_silh].sum() / batch_size | |
| else: | |
| print('silh_loss_type: ' + silh_loss_type) | |
| raise ValueError | |
| # shape regularization | |
| # 'smal': loss on betas (pca coefficients), betas should be close to 0 | |
| # 'limbs...' loss on selected betas_limbs | |
| loss_shape_weighted_list = [torch.zeros((1)).mean().to(output_reproj['keyp_2d'].device)] | |
| for ind_sp, sp in enumerate(weight_dict['shape_options']): | |
| weight_sp = weight_dict['shape'][ind_sp] | |
| # self.logscale_part_list = ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'] | |
| if sp == 'smal': | |
| loss_shape_tmp = self.shape_prior(output_reproj['betas']) | |
| elif sp == 'limbs': | |
| loss_shape_tmp = torch.mean((output_reproj['betas_limbs'])**2) | |
| elif sp == 'limbs7': | |
| limb_coeffs_list = [0.01, 1, 0.1, 1, 1, 0.1, 2] | |
| limb_coeffs = torch.tensor(limb_coeffs_list).to(torch.float32).to(target_dict['tpts'].device) | |
| loss_shape_tmp = torch.mean((output_reproj['betas_limbs'] * limb_coeffs[None, :])**2) | |
| else: | |
| raise NotImplementedError | |
| loss_shape_weighted_list.append(weight_sp * loss_shape_tmp) | |
| loss_shape_weighted = torch.stack((loss_shape_weighted_list)).sum() | |
| # 3D loss for dogs for which we have a unity model or toy figure | |
| loss_models3d = torch.zeros((1)).mean().to(output_reproj['betas'].device) | |
| if 'models3d' in weight_dict.keys(): | |
| if weight_dict['models3d'] > 0: | |
| for ind_dog in range(target_dict['breed_index'].shape[0]): | |
| breed_index = np.asscalar(target_dict['breed_index'][ind_dog].detach().cpu().numpy()) | |
| if breed_index in self.dog_betas_unity.keys(): | |
| betas_target = self.dog_betas_unity[breed_index][:output_reproj['betas'].shape[1]].to(output_reproj['betas'].device) | |
| betas_output = output_reproj['betas'][ind_dog, :] | |
| betas_limbs_output = output_reproj['betas_limbs'][ind_dog, :] | |
| loss_models3d += ((betas_limbs_output**2).sum() + ((betas_output-betas_target)**2).sum()) / (output_reproj['betas'].shape[1] + output_reproj['betas_limbs'].shape[1]) | |
| else: | |
| weight_dict['models3d'] = 0 | |
| # shape resularization loss on shapedirs | |
| # -> in the current version shapedirs are kept fixed, so we don't need those losses | |
| if weight_dict['shapedirs'] > 0: | |
| raise NotImplementedError | |
| else: | |
| loss_shapedirs = torch.zeros((1)).mean().to(output_reproj['betas'].device) | |
| # prior on back joints (not used in cvpr 2022 paper) | |
| # -> elementwise MSE loss on all 6 coefficients of 6d rotation representation | |
| if 'pose_0' in weight_dict.keys(): | |
| if weight_dict['pose_0'] > 0: | |
| pred_pose_rot6d = output_reproj['pose_rot6d'] | |
| w_rj_np = np.zeros((pred_pose_rot6d.shape[1])) | |
| w_rj_np[[2, 3, 4, 5]] = 1.0 # back | |
| w_rj = torch.tensor(w_rj_np).to(torch.float32).to(pred_pose_rot6d.device) | |
| zero_rot = torch.tensor([1, 0, 0, 1, 0, 0]).to(pred_pose_rot6d.device).to(torch.float32)[None, None, :].repeat((batch_size, pred_pose_rot6d.shape[1], 1)) | |
| loss_pose = self.criterion_regr(pred_pose_rot6d*w_rj[None, :, None], zero_rot*w_rj[None, :, None]) | |
| else: | |
| loss_pose = torch.zeros((1)).mean() | |
| # pose prior | |
| # -> we did experiment with different pose priors, for example: | |
| # * similart to SMALify (https://github.com/benjiebob/SMALify/blob/master/smal_fitter/smal_fitter.py, | |
| # https://github.com/benjiebob/SMALify/blob/master/smal_fitter/priors/pose_prior_35.py) | |
| # * vae | |
| # * normalizing flow pose prior | |
| # -> our cvpr 2022 paper uses the normalizing flow pose prior as implemented below | |
| if 'poseprior' in weight_dict.keys(): | |
| if weight_dict['poseprior'] > 0: | |
| pred_pose_rot6d = output_reproj['pose_rot6d'] | |
| pred_pose = rot6d_to_rotmat(pred_pose_rot6d.reshape((-1, 6))).reshape((batch_size, -1, 3, 3)) | |
| if 'normalizing_flow_tiger' in weight_dict['poseprior_options']: | |
| if output_reproj['normflow_z'] is not None: | |
| loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss_from_z(output_reproj['normflow_z'], type='square') | |
| else: | |
| loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss(pred_pose_rot6d, type='square') | |
| elif 'normalizing_flow_tiger_logprob' in weight_dict['poseprior_options']: | |
| if output_reproj['normflow_z'] is not None: | |
| loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss_from_z(output_reproj['normflow_z'], type='neg_log_prob') | |
| else: | |
| loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss(pred_pose_rot6d, type='neg_log_prob') | |
| else: | |
| raise NotImplementedError | |
| else: | |
| loss_poseprior = torch.zeros((1)).mean() | |
| else: | |
| weight_dict['poseprior'] = 0 | |
| loss_poseprior = torch.zeros((1)).mean() | |
| # add a prior which penalizes side-movement angles for legs | |
| if 'poselegssidemovement' in weight_dict.keys(): | |
| use_pose_legs_side_loss = True | |
| else: | |
| use_pose_legs_side_loss = False | |
| if use_pose_legs_side_loss: | |
| leg_indices_right = np.asarray([7, 8, 9, 10, 17, 18, 19, 20]) # front, back | |
| leg_indices_left = np.asarray([11, 12, 13, 14, 21, 22, 23, 24]) # front, back | |
| vec = torch.zeros((3, 1)).to(device=pred_pose.device, dtype=pred_pose.dtype) | |
| vec[2] = -1 | |
| x0_rotmat = pred_pose | |
| x0_rotmat_legs_left = x0_rotmat[:, leg_indices_left, :, :] | |
| x0_rotmat_legs_right = x0_rotmat[:, leg_indices_right, :, :] | |
| x0_legs_left = x0_rotmat_legs_left.reshape((-1, 3, 3))@vec | |
| x0_legs_right = x0_rotmat_legs_right.reshape((-1, 3, 3))@vec | |
| eps=0 # 1e-7 | |
| # use the component of the vector which points to the side | |
| loss_poselegssidemovement = (x0_legs_left[:, 1]**2).mean() + (x0_legs_right[:, 1]**2).mean() | |
| else: | |
| loss_poselegssidemovement = torch.zeros((1)).mean() | |
| weight_dict['poselegssidemovement'] = 0 | |
| # dog breed classification loss | |
| dog_breed_gt = target_dict['breed_index'] | |
| dog_breed_pred = output_reproj['dog_breed'] | |
| loss_class = self.criterion_class(dog_breed_pred, dog_breed_gt) | |
| # dog breed relationship loss | |
| # -> we did experiment with many other options, but none was significantly better | |
| if '4' in weight_dict['breed_options']: # we have pairs of dogs of the same breed | |
| assert weight_dict['breed'] > 0 | |
| z = output_reproj['z'] | |
| # go through all pairs and compare them to each other sample | |
| if self.l_anchor is None: | |
| self.prepare_anchor_pos_neg(batch_size, z.device) | |
| anchor = torch.index_select(z, 0, self.l_anchor) | |
| positive = torch.index_select(z, 0, self.l_pos) | |
| negative = torch.index_select(z, 0, self.l_neg) | |
| loss_breed = self.criterion_triplet(anchor, positive, negative) | |
| else: | |
| loss_breed = torch.zeros((1)).mean() | |
| # regularizarion for focal length | |
| loss_flength_near_mean = torch.mean(output_reproj['flength']**2) | |
| loss_flength = loss_flength_near_mean | |
| # bodypart segmentation loss | |
| if 'partseg' in weight_dict.keys(): | |
| if weight_dict['partseg'] > 0: | |
| raise NotImplementedError | |
| else: | |
| loss_partseg = torch.zeros((1)).mean() | |
| else: | |
| weight_dict['partseg'] = 0 | |
| loss_partseg = torch.zeros((1)).mean() | |
| # weight and combine losses | |
| loss_keyp_weighted = loss_keyp * weight_dict['keyp'] | |
| loss_silh_weighted = loss_silh * weight_dict['silh'] | |
| loss_shapedirs_weighted = loss_shapedirs * weight_dict['shapedirs'] | |
| loss_pose_weighted = loss_pose * weight_dict['pose_0'] | |
| loss_class_weighted = loss_class * weight_dict['class'] | |
| loss_breed_weighted = loss_breed * weight_dict['breed'] | |
| loss_flength_weighted = loss_flength * weight_dict['flength'] | |
| loss_poseprior_weighted = loss_poseprior * weight_dict['poseprior'] | |
| loss_partseg_weighted = loss_partseg * weight_dict['partseg'] | |
| loss_models3d_weighted = loss_models3d * weight_dict['models3d'] | |
| loss_poselegssidemovement_weighted = loss_poselegssidemovement * weight_dict['poselegssidemovement'] | |
| #################################################################################################### | |
| loss = loss_keyp_weighted + loss_silh_weighted + loss_shape_weighted + loss_pose_weighted + loss_class_weighted + \ | |
| loss_shapedirs_weighted + loss_breed_weighted + loss_flength_weighted + loss_poseprior_weighted + \ | |
| loss_partseg_weighted + loss_models3d_weighted + loss_poselegssidemovement_weighted | |
| #################################################################################################### | |
| loss_dict = {'loss': loss.item(), | |
| 'loss_keyp_weighted': loss_keyp_weighted.item(), \ | |
| 'loss_silh_weighted': loss_silh_weighted.item(), \ | |
| 'loss_shape_weighted': loss_shape_weighted.item(), \ | |
| 'loss_shapedirs_weighted': loss_shapedirs_weighted.item(), \ | |
| 'loss_pose0_weighted': loss_pose_weighted.item(), \ | |
| 'loss_class_weighted': loss_class_weighted.item(), \ | |
| 'loss_breed_weighted': loss_breed_weighted.item(), \ | |
| 'loss_flength_weighted': loss_flength_weighted.item(), \ | |
| 'loss_poseprior_weighted': loss_poseprior_weighted.item(), \ | |
| 'loss_partseg_weighted': loss_partseg_weighted.item(), \ | |
| 'loss_models3d_weighted': loss_models3d_weighted.item(), \ | |
| 'loss_poselegssidemovement_weighted': loss_poselegssidemovement_weighted.item()} | |
| return loss, loss_dict | |