barc_gradio / src /combined_model /loss_image_to_3d_withbreedrel.py
Nadine Rueegg
initial commit for barc
7629b39
raw
history blame contribute delete
No virus
15.6 kB
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