code: """
  Loss Utils.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import torch
  from . import geom_utils
  import numpy as np
  
  def mask_dt_loss(proj_verts, dist_transf):
      """
      proj_verts: B x N x 2
      (In normalized coordinate [-1, 1])
      dist_transf: B x 1 x N x N
  
      Computes the distance transform at the points where vertices land.
      """
      # Reshape into B x 1 x N x 2
      sample_grid = proj_verts.unsqueeze(1)
      # B x 1 x 1 x N
      dist_transf = torch.nn.functional.grid_sample(dist_transf, sample_grid, padding_mode='border')
      return dist_transf.mean()
  
  
  def texture_dt_loss(texture_flow, dist_transf, vis_rend=None, cams=None, verts=None, tex_pred=None):
      """
      texture_flow: B x F x T x T x 2
      (In normalized coordinate [-1, 1])
      dist_transf: B x 1 x N x N
  
      Similar to geom_utils.sample_textures
      But instead of sampling image, it samples dt values.
      """
      # Reshape into B x F x T*T x 2
      T = texture_flow.size(-2)
      F = texture_flow.size(1)
      flow_grid = texture_flow.view(-1, F, T * T, 2)
      # B x 1 x F x T*T
      dist_transf = torch.nn.functional.grid_sample(dist_transf, flow_grid)
  
      if vis_rend is not None:
          # Visualize the error!
          # B x 3 x F x T*T
          dts = dist_transf.repeat(1, 3, 1, 1)
          # B x 3 x F x T x T
          dts = dts.view(-1, 3, F, T, T)
          # B x F x T x T x 3
          dts = dts.permute(0, 2, 3, 4, 1)
          dts = dts.unsqueeze(4).repeat(1, 1, 1, 1, T, 1) / dts.max()
  
          from ..utils import bird_vis
          for i in range(dist_transf.size(0)):
              rend_dt = vis_rend(verts[i], cams[i], dts[i])
              rend_img = bird_vis.tensor2im(tex_pred[i].data)            
              import matplotlib.pyplot as plt
              plt.ion()
              fig=plt.figure(1)
              plt.clf()
              ax = fig.add_subplot(121)
              ax.imshow(rend_dt)
              ax = fig.add_subplot(122)
              ax.imshow(rend_img)
              import ipdb; ipdb.set_trace()
  
      return dist_transf.mean()
  
  
  def texture_loss(img_pred, img_gt, mask_pred, mask_gt):
      """
      Input:
        img_pred, img_gt: B x 3 x H x W
        mask_pred, mask_gt: B x H x W
      """
      mask_pred = mask_pred.unsqueeze(1)
      mask_gt = mask_gt.unsqueeze(1)
  
      # masked_rend = (img_pred * mask)[0].data.cpu().numpy()
      # masked_gt = (img_gt * mask)[0].data.cpu().numpy()
      # import matplotlib.pyplot as plt
      # plt.ion()
      # plt.figure(1)
      # plt.clf()
      # fig = plt.figure(1)
      # ax = fig.add_subplot(121)
      # ax.imshow(np.transpose(masked_rend, (1, 2, 0)))
      # ax = fig.add_subplot(122)
      # ax.imshow(np.transpose(masked_gt, (1, 2, 0)))
      # import ipdb; ipdb.set_trace()
  
      return torch.nn.L1Loss()(img_pred * mask_pred, img_gt * mask_gt)
  
  
  def camera_loss(cam_pred, cam_gt, margin):
      """
      cam_* are B x 7, [sc, tx, ty, quat]
      Losses are in similar magnitude so one margin is ok.
      """
      # CH: camera loss always 0
      return 0
  
      # CH: comment out old loss code
  #     rot_pred = cam_pred[:, -4:]
  #     rot_gt = cam_gt[:, -4:]
  
  #     rot_loss = hinge_loss(quat_loss_geodesic(rot_pred, rot_gt), margin)
  #     # Scale and trans.
  #     st_loss = (cam_pred[:, :3] - cam_gt[:, :3])**2
  #     st_loss = hinge_loss(st_loss.view(-1), margin)
  
  #     return rot_loss.mean() + st_loss.mean()
  
  def hinge_loss(loss, margin):
      # Only penalize if loss > margin
      zeros = torch.autograd.Variable(torch.zeros(1).cuda(), requires_grad=False)
      return torch.max(loss - margin, zeros)
  
  
  def quat_loss_geodesic(q1, q2):
      '''
      Geodesic rotation loss.
      
      Args:
          q1: N X 4
          q2: N X 4
      Returns:
          loss : N x 1
      '''
      q1 = torch.unsqueeze(q1, 1)
      q2 = torch.unsqueeze(q2, 1)
      q2_conj = torch.cat([ q2[:, :, [0]] , -1*q2[:, :, 1:4] ], dim=-1)
      q_rel = geom_utils.hamilton_product(q1, q2_conj)
      q_loss = 1 - torch.abs(q_rel[:, :, 0])
      # we can also return q_loss*q_loss
      return q_loss
      
  
  def quat_loss(q1, q2):
      '''
      Anti-podal squared L2 loss.
      
      Args:
          q1: N X 4
          q2: N X 4
      Returns:
          loss : N x 1
      '''
      q_diff_loss = (q1-q2).pow(2).sum(1)
      q_sum_loss = (q1+q2).pow(2).sum(1)
      q_loss, _ = torch.stack((q_diff_loss, q_sum_loss), dim=1).min(1)
      return q_loss
  
  
  def triangle_loss(verts, edge2verts):
      """
      Encourages dihedral angle to be 180 degrees.
  
      Args:
          verts: B X N X 3
          edge2verts: B X E X 4
      Returns:
          loss : scalar
      """
      indices_repeat = torch.stack([edge2verts, edge2verts, edge2verts], dim=2) # B X E X 3 X 4
  
      verts_A = torch.gather(verts, 1, indices_repeat[:, :, :, 0])
      verts_B = torch.gather(verts, 1, indices_repeat[:, :, :, 1])
      verts_C = torch.gather(verts, 1, indices_repeat[:, :, :, 2])
      verts_D = torch.gather(verts, 1, indices_repeat[:, :, :, 3])
  
      # n1 = cross(ad, ab)
      # n2 = cross(ab, ac)
      n1 = geom_utils.cross_product(verts_D - verts_A, verts_B - verts_A)
      n2 = geom_utils.cross_product(verts_B - verts_A, verts_C - verts_A)
  
      n1 = torch.nn.functional.normalize(n1, dim=2)
      n2 = torch.nn.functional.normalize(n2, dim=2)
  
      dot_p = (n1 * n2).sum(2)
      loss = ((1 - dot_p)**2).mean()
      return loss
  
  
  def deform_l2reg(V):
      """
      l2 norm on V = B x N x 3
      """
      V = V.view(-1, V.size(2))
      return torch.mean(torch.norm(V, p=2, dim=1))
  
  
  def entropy_loss(A):
      """
      Input is K x N
      Each column is a prob of vertices being the one for k-th keypoint.
      We want this to be sparse = low entropy.
      """
      entropy = -torch.sum(A * torch.log(A), 1)
      # Return avg entropy over 
      return torch.mean(entropy)
  
  
  def kp_l2_loss(kp_pred, kp_gt):
      """
      L2 loss between visible keypoints.
  
      \Sum_i [0.5 * vis[i] * (kp_gt[i] - kp_pred[i])^2] / (|vis|)
      """
      # CH: Do not consider keypoint loss
      return 0
  
  #     criterion = torch.nn.MSELoss()
  
  #     vis = (kp_gt[:, :, 2, None] > 0).float()
  
  #     # This always has to be (output, target), not (target, output)
  #     return criterion(vis * kp_pred, vis * kp_gt[:, :, :2])
  
  
  def lsgan_loss(score_real, score_fake):
      """
      DELETE ME.
      Label 0=fake, 1=real.
      score_real is B x 1, score for real samples
      score_fake is B x 1, score for fake samples
  
      Returns loss for discriminator and encoder.
      """
  
      disc_loss_real = torch.mean((score_real - 1)**2)
      disc_loss_fake = torch.mean((score_fake)**2)
      disc_loss = disc_loss_real + disc_loss_fake
  
      enc_loss = torch.mean((score_fake - 1)**2)
  
      return disc_loss, enc_loss
  
  
  class EdgeLoss(object):
      """
      Edge length should not diverge from the original edge length.
  
      On initialization computes the current edge lengths.
      """
      def __init__(self, verts, edges2verts, margin=2, use_bad_edge=False, use_l2=False):
          # Input:
          #  verts: B x N x 3
          #  edeges2verts: B x E x 4
          #  (only using the first 2 columns)
          self.use_l2 = use_l2
  
          # B x E x 2
          edge_verts = edges2verts[:, :, :2]
          self.indices = torch.stack([edge_verts, edge_verts, edge_verts], dim=2)
          V_copy = torch.autograd.Variable(verts.data, requires_grad=False)
          if V_copy.dim() == 2:
              # N x 3 (mean shape) -> B x N x 3
              V_copy = V_copy.unsqueeze(0).repeat(edges2verts.size(0), 1, 1)
  
          if use_bad_edge:
              self.log_e0 = torch.log(self.compute_edgelength(V_copy))
          else:
              # e0 is the mean over all edge lengths!
              e0 = self.compute_edgelength(V_copy).mean(1).view(-1, 1)
              self.log_e0 = torch.log(e0)
  
          self.margin = np.log(margin)
          self.zeros = torch.autograd.Variable(torch.zeros(1).cuda(), requires_grad=False)
  
          # For visualization
          self.v1 = edges2verts[0, :, 0].data.cpu().numpy()
          self.v2 = edges2verts[0, :, 1].data.cpu().numpy()
  
      def __call__(self, verts):
          e1 = self.compute_edgelength(verts)
          if self.use_l2:
              dist = (torch.log(e1) - self.log_e0)**2
              self.dist = torch.max(dist - self.margin**2, self.zeros)
          else:
              dist = torch.abs(torch.log(e1) - self.log_e0)
              self.dist = torch.max(dist - self.margin, self.zeros)
          return self.dist.mean()
  
      def compute_edgelength(self, V):
          v1 = torch.gather(V, 1, self.indices[:, :, :, 0])
          v2 = torch.gather(V, 1, self.indices[:, :, :, 1])
  
          elengths = torch.sqrt(((v1 - v2)**2).sum(2))
  
          # B x E
          return elengths
  
      def visualize(self, verts, F_np, mv=None):
          from psbody.mesh import Mesh
  
          V = verts[0].data.cpu().numpy()
          mesh = Mesh(V, F_np)
          dist = self.dist[0].data.cpu().numpy()
  
          v_weights = np.zeros((V.shape[0]))
          for e_id, (v1_id, v2_id) in enumerate(zip(self.v1, self.v2)):
              v_weights[v1_id] += dist[e_id]
              v_weights[v2_id] += dist[e_id]
  
          mesh.set_vertex_colors_from_weights(v_weights)
  
          if mv is not None:
              mv.set_dynamic_meshes([mesh])
          else:
              mesh.show()
              import ipdb; ipdb.set_trace()
  
  
  class LaplacianLoss(object):
      """
      Encourages minimal mean curvature shapes.
      """
      def __init__(self, faces):
          # Input:
          #  faces: B x F x 3
          from ..nnutils.laplacian import Laplacian
          # V x V
          self.laplacian = Laplacian(faces)
          self.Lx = None
  
      def __call__(self, verts):
          self.Lx = self.laplacian(verts)
          # Reshape to BV x 3
          Lx = self.Lx.view(-1, self.Lx.size(2))
          loss = torch.norm(Lx, p=2, dim=1).mean()
          return loss
  
      def visualize(self, verts, mv=None):
          # Visualizes the laplacian.
          # Verts is B x N x 3 Variable
          Lx = self.Lx[0].data.cpu().numpy()
  
          V = verts[0].data.cpu().numpy()
  
          from psbody.mesh import Mesh
          F = self.laplacian.F_np[0]
          mesh = Mesh(V, F)
  
          weights = np.linalg.norm(Lx, axis=1)
          mesh.set_vertex_colors_from_weights(weights)
  
          if mv is not None:
              mv.set_dynamic_meshes([mesh])
          else:
              mesh.show()
              import ipdb; ipdb.set_trace()
  
  
  class PerceptualTextureLoss(object):
      def __init__(self):
          from ..nnutils.perceptual_loss import PerceptualLoss
          self.perceptual_loss = PerceptualLoss()
  
      def __call__(self, img_pred, img_gt, mask_pred, mask_gt):
          """
          Input:
            img_pred, img_gt: B x 3 x H x W
          mask_pred, mask_gt: B x H x W
          """
          mask_pred = mask_pred.unsqueeze(1)
          mask_gt = mask_gt.unsqueeze(1)
          # masked_rend = (img_pred * mask_pred)[0].data.cpu().numpy()
          # masked_gt = (img_gt * mask_gt)[0].data.cpu().numpy()
          # import matplotlib.pyplot as plt
          # plt.ion()
          # plt.figure(1)
          # plt.clf()
          # fig = plt.figure(1)
          # ax = fig.add_subplot(121)
          # ax.imshow(np.transpose(masked_rend, (1, 2, 0)))
          # ax = fig.add_subplot(122)
          # ax.imshow(np.transpose(masked_gt, (1, 2, 0)))
          # import ipdb; ipdb.set_trace()
  
          # Only use mask_gt..
          dist = self.perceptual_loss(img_pred * mask_gt, img_gt * mask_gt)
          return dist.mean()
  
apis: torch.stack torch.nn.L1Loss numpy.log torch.log torch.nn.functional.grid_sample torch.max torch.cat matplotlib.pyplot.figure torch.autograd.Variable torch.gather torch.norm torch.mean torch.unsqueeze numpy.zeros torch.nn.functional.normalize matplotlib.pyplot.clf matplotlib.pyplot.ion numpy.linalg.norm torch.abs torch.zeros