RabbitRUI/ruispace

1import numpy as np2import torch3import torch.nn as nn4from kornia.geometry import warp_affine5import torch.nn.functional as F6 7def resize_n_crop(image, M, dsize=112):8    # image: (b, c, h, w)9    # M   :  (b, 2, 3)10    return warp_affine(image, M, dsize=(dsize, dsize), align_corners=True)11 12### perceptual level loss13class PerceptualLoss(nn.Module):14    def __init__(self, recog_net, input_size=112):15        super(PerceptualLoss, self).__init__()16        self.recog_net = recog_net17        self.preprocess = lambda x: 2 * x - 118        self.input_size=input_size19    def forward(imageA, imageB, M):20        """21        1 - cosine distance22        Parameters:23            imageA       --torch.tensor (B, 3, H, W), range (0, 1) , RGB order24            imageB       --same as imageA25        """26 27        imageA = self.preprocess(resize_n_crop(imageA, M, self.input_size))28        imageB = self.preprocess(resize_n_crop(imageB, M, self.input_size))29 30        # freeze bn31        self.recog_net.eval()32        33        id_featureA = F.normalize(self.recog_net(imageA), dim=-1, p=2)34        id_featureB = F.normalize(self.recog_net(imageB), dim=-1, p=2)  35        cosine_d = torch.sum(id_featureA * id_featureB, dim=-1)36        # assert torch.sum((cosine_d > 1).float()) == 037        return torch.sum(1 - cosine_d) / cosine_d.shape[0]        38 39def perceptual_loss(id_featureA, id_featureB):40    cosine_d = torch.sum(id_featureA * id_featureB, dim=-1)41        # assert torch.sum((cosine_d > 1).float()) == 042    return torch.sum(1 - cosine_d) / cosine_d.shape[0]  43 44### image level loss45def photo_loss(imageA, imageB, mask, eps=1e-6):46    """47    l2 norm (with sqrt, to ensure backward stabililty, use eps, otherwise Nan may occur)48    Parameters:49        imageA       --torch.tensor (B, 3, H, W), range (0, 1), RGB order 50        imageB       --same as imageA51    """52    loss = torch.sqrt(eps + torch.sum((imageA - imageB) ** 2, dim=1, keepdims=True)) * mask53    loss = torch.sum(loss) / torch.max(torch.sum(mask), torch.tensor(1.0).to(mask.device))54    return loss55 56def landmark_loss(predict_lm, gt_lm, weight=None):57    """58    weighted mse loss59    Parameters:60        predict_lm    --torch.tensor (B, 68, 2)61        gt_lm         --torch.tensor (B, 68, 2)62        weight        --numpy.array (1, 68)63    """64    if not weight:65        weight = np.ones([68])66        weight[28:31] = 2067        weight[-8:] = 2068        weight = np.expand_dims(weight, 0)69        weight = torch.tensor(weight).to(predict_lm.device)70    loss = torch.sum((predict_lm - gt_lm)**2, dim=-1) * weight71    loss = torch.sum(loss) / (predict_lm.shape[0] * predict_lm.shape[1])72    return loss73 74 75### regulization76def reg_loss(coeffs_dict, opt=None):77    """78    l2 norm without the sqrt, from yu's implementation (mse)79    tf.nn.l2_loss https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss80    Parameters:81        coeffs_dict     -- a  dict of torch.tensors , keys: id, exp, tex, angle, gamma, trans82 83    """84    # coefficient regularization to ensure plausible 3d faces85    if opt:86        w_id, w_exp, w_tex = opt.w_id, opt.w_exp, opt.w_tex87    else:88        w_id, w_exp, w_tex = 1, 1, 1, 189    creg_loss = w_id * torch.sum(coeffs_dict['id'] ** 2) +  \90           w_exp * torch.sum(coeffs_dict['exp'] ** 2) + \91           w_tex * torch.sum(coeffs_dict['tex'] ** 2)92    creg_loss = creg_loss / coeffs_dict['id'].shape[0]93 94    # gamma regularization to ensure a nearly-monochromatic light95    gamma = coeffs_dict['gamma'].reshape([-1, 3, 9])96    gamma_mean = torch.mean(gamma, dim=1, keepdims=True)97    gamma_loss = torch.mean((gamma - gamma_mean) ** 2)98 99    return creg_loss, gamma_loss100 101def reflectance_loss(texture, mask):102    """103    minimize texture variance (mse), albedo regularization to ensure an uniform skin albedo104    Parameters:105        texture       --torch.tensor, (B, N, 3)106        mask          --torch.tensor, (N), 1 or 0107 108    """109    mask = mask.reshape([1, mask.shape[0], 1])110    texture_mean = torch.sum(mask * texture, dim=1, keepdims=True) / torch.sum(mask)111    loss = torch.sum(((texture - texture_mean) * mask)**2) / (texture.shape[0] * torch.sum(mask))112    return loss113 114