text: import torch
  from torch.autograd import Variable
  
  
  class PixelwiseContrastiveLoss(object):
  
      def __init__(self, image_shape, config=None):
          self.type = "pixelwise_contrastive"
          self.image_width  = image_shape[1]
          self.image_height = image_shape[0]
  
          assert config is not None
          self._config = config
  
          self._debug_data = dict()
  
          self._debug = False
  
      @property
      def debug(self):
          return self._debug
  
      @property
      def config(self):
          return self._config
  
      @debug.setter
      def debug(self, value):
          self._debug = value
  
      @property
      def debug_data(self):
          return self._debug_data
  
      def get_loss_matched_and_non_matched_with_l2(self, image_a_pred, image_b_pred, matches_a, matches_b, non_matches_a, non_matches_b,
                   M_descriptor=None, M_pixel=None, non_match_loss_weight=1.0, use_l2_pixel_loss=None):
          """
          Computes the loss function
  
          DCN = Dense Correspondence Network
          num_images = number of images in this batch
          num_matches = number of matches
          num_non_matches = number of non-matches
          W = image width
          H = image height
          D = descriptor dimension
  
  
          match_loss = 1/num_matches \sum_{num_matches} ||descriptor_a - descriptor_b||_2^2
          non_match_loss = 1/num_non_matches \sum_{num_non_matches} max(0, M_margin - ||descriptor_a - descriptor_b||_2)^2
  
          loss = match_loss + non_match_loss
  
          :param image_a_pred: Output of DCN network on image A.
          :type image_a_pred: torch.Variable(torch.FloatTensor) shape [1, W * H, D]
          :param image_b_pred: same as image_a_pred
          :type image_b_pred:
          :param matches_a: torch.Variable(torch.LongTensor) has shape [num_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of one dimension of image_a_pred
          :type matches_a: torch.Variable(torch.FloatTensor)
          :param matches_b: same as matches_a
          :type matches_b:
          :param non_matches_a: torch.Variable(torch.FloatTensor) has shape [num_non_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of image_a_pred
          :type non_matches_a: torch.Variable(torch.FloatTensor)
          :param non_matches_b: same as non_matches_a
          :type non_matches_b:
          :return: loss, match_loss, non_match_loss
          :rtype: torch.Variable(torch.FloatTensor) each of shape torch.Size([1])
          """
  
          PCL = PixelwiseContrastiveLoss
  
          if M_descriptor is None:
              M_descriptor = self._config["M_descriptor"]
  
          if M_pixel is None:
              M_pixel = self._config["M_pixel"]
  
  
          if use_l2_pixel_loss is None:
              use_l2_pixel_loss = self._config['use_l2_pixel_loss_on_masked_non_matches']
  
  
          match_loss, _, _ = PCL.match_loss(image_a_pred, image_b_pred, matches_a, matches_b)
  
  
  
          if use_l2_pixel_loss:
              non_match_loss, num_hard_negatives =\
                  self.non_match_loss_with_l2_pixel_norm(image_a_pred, image_b_pred, matches_b,
                                                         non_matches_a, non_matches_b,
                                                         M_descriptor=M_descriptor,
                                                         M_pixel=M_pixel)
          else:
              # version with no l2 pixel term
              non_match_loss, num_hard_negatives = self.non_match_loss_descriptor_only(image_a_pred, image_b_pred, non_matches_a, non_matches_b, M_descriptor=M_descriptor)
  
  
  
          return match_loss, non_match_loss, num_hard_negatives
  
      @staticmethod
      def get_triplet_loss(image_a_pred, image_b_pred, matches_a, matches_b, non_matches_a, non_matches_b, alpha):
          """
          Computes the loss function
  
          \sum_{triplets} ||D(I_a, u_a, I_b, u_{b,match})||_2^2 - ||D(I_a, u_a, I_b, u_{b,non-match)||_2^2 + alpha 
  
          """
          num_matches = matches_a.size()[0]
          num_non_matches = non_matches_a.size()[0]
          multiplier = num_non_matches / num_matches
  
          ## non_matches_a is already replicated up to be the right size
          ## non_matches_b is also that side
          ## matches_a is just a smaller version of non_matches_a
          ## matches_b is the only thing that needs to be replicated up in size
  
          matches_b_long =  torch.t(matches_b.repeat(multiplier, 1)).contiguous().view(-1)
                           
          matches_a_descriptors = torch.index_select(image_a_pred, 1, non_matches_a)
          matches_b_descriptors      = torch.index_select(image_b_pred, 1, matches_b_long)
          non_matches_b_descriptors  = torch.index_select(image_b_pred, 1, non_matches_b)
  
          triplet_losses = (matches_a_descriptors - matches_b_descriptors).pow(2) - (matches_a_descriptors - non_matches_b_descriptors).pow(2) + alpha
          triplet_loss = 1.0 / num_non_matches * torch.clamp(triplet_losses, min=0).sum()
  
          return triplet_loss
  
      @staticmethod
      def match_loss(image_a_pred, image_b_pred, matches_a, matches_b):
          """
          Computes the match loss given by
  
          1/num_matches * \sum_{matches} ||D(I_a, u_a, I_b, u_b)||_2^2
  
          :param image_a_pred: Output of DCN network on image A.
          :type image_a_pred: torch.Variable(torch.FloatTensor) shape [1, W * H, D]
          :param image_b_pred: same as image_a_pred
          :type image_b_pred:
          :param matches_a: torch.Variable(torch.LongTensor) has shape [num_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of one dimension of image_a_pred
          :type matches_a: torch.Variable(torch.FloatTensor)
          :param matches_b: same as matches_b
  
          :return: match_loss, matches_a_descriptors, matches_b_descriptors
          :rtype: torch.Variable(),
  
          matches_a_descriptors is torch.FloatTensor with shape torch.Shape([num_matches, descriptor_dimension])
          """
  
          num_matches = matches_a.size()[0]
          matches_a_descriptors = torch.index_select(image_a_pred, 1, matches_a)
          matches_b_descriptors = torch.index_select(image_b_pred, 1, matches_b)
  
          # crazily enough, if there is only one element to index_select into
          # above, then the first dimension is collapsed down, and we end up 
          # with shape [D,], where we want [1,D]
          # this unsqueeze fixes that case
          if len(matches_a) == 1:
              matches_a_descriptors = matches_a_descriptors.unsqueeze(0)
              matches_b_descriptors = matches_b_descriptors.unsqueeze(0)
  
          if num_matches == 0:
              match_loss = torch.zeros((1)).cuda()
          else:
              match_loss = 1.0 / num_matches * (matches_a_descriptors - matches_b_descriptors).pow(2).sum()
  
          return match_loss, matches_a_descriptors, matches_b_descriptors
  
  
      @staticmethod
      def non_match_descriptor_loss(image_a_pred, image_b_pred, non_matches_a, non_matches_b, M=0.5, invert=False):
          """
          Computes the max(0, M - D(I_a,I_b,u_a,u_b))^2 term
  
          This is effectively:       "a and b should be AT LEAST M away from each other"
          With invert=True, this is: "a and b should be AT MOST  M away from each other" 
  
           :param image_a_pred: Output of DCN network on image A.
          :type image_a_pred: torch.Variable(torch.FloatTensor) shape [1, W * H, D]
          :param image_b_pred: same as image_a_pred
          :type image_b_pred:
          :param non_matches_a: torch.Variable(torch.FloatTensor) has shape [num_non_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of image_a_pred
          :type non_matches_a: torch.Variable(torch.FloatTensor)
          :param non_matches_b: same as non_matches_a
          :param M: the margin
          :type M: float
          :return: torch.FloatTensor with shape torch.Shape([num_non_matches])
          :rtype:
          """
  
          non_matches_a_descriptors = torch.index_select(image_a_pred, 1, non_matches_a).squeeze()
          non_matches_b_descriptors = torch.index_select(image_b_pred, 1, non_matches_b).squeeze()
  
          # crazily enough, if there is only one element to index_select into
          # above, then the first dimension is collapsed down, and we end up 
          # with shape [D,], where we want [1,D]
          # this unsqueeze fixes that case
          if len(non_matches_a) == 1:
              non_matches_a_descriptors = non_matches_a_descriptors.unsqueeze(0)
              non_matches_b_descriptors = non_matches_b_descriptors.unsqueeze(0)
  
          norm_degree = 2
          non_match_loss = (non_matches_a_descriptors - non_matches_b_descriptors).norm(norm_degree, 1)
          if not invert:
              non_match_loss = torch.clamp(M - non_match_loss, min=0).pow(2)
          else:
              non_match_loss = torch.clamp(non_match_loss - M, min=0).pow(2)
  
          hard_negative_idxs = torch.nonzero(non_match_loss)
          num_hard_negatives = len(hard_negative_idxs)
  
          return non_match_loss, num_hard_negatives, non_matches_a_descriptors, non_matches_b_descriptors
  
      def non_match_loss_with_l2_pixel_norm(self, image_a_pred, image_b_pred, matches_b,
                                            non_matches_a, non_matches_b, M_descriptor=0.5,
                                            M_pixel=None):
  
          """
  
          Computes the total non_match_loss with an l2_pixel norm term
  
          :param image_a_pred: Output of DCN network on image A.
          :type image_a_pred: torch.Variable(torch.FloatTensor) shape [1, W * H, D]
          :param image_b_pred: same as image_a_pred
          :type image_b_pred:
          :param matches_a: torch.Variable(torch.LongTensor) has shape [num_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of one dimension of image_a_pred
          :type matches_a: torch.Variable(torch.FloatTensor)
          :param matches_b: same as matches_a
          :type matches_b:
          :param non_matches_a: torch.Variable(torch.FloatTensor) has shape [num_non_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of image_a_pred
          :type non_matches_a: torch.Variable(torch.FloatTensor)
          :param non_matches_b: same as non_matches_a
  
          :param M_descriptor: margin for descriptor loss term
          :type M_descriptor: float
          :param M_pixel: margin for pixel loss term
          :type M_pixel: float
          :return: non_match_loss, num_hard_negatives
          :rtype: torch.Variable, int
          """
  
          if M_descriptor is None:
              M_descriptor = self._config["M_descriptor"]
  
          if M_pixel is None:
              M_pixel = self._config["M_pixel"]
  
          PCL = PixelwiseContrastiveLoss
  
          num_non_matches = non_matches_a.size()[0]
  
          non_match_descriptor_loss, num_hard_negatives, _, _ = PCL.non_match_descriptor_loss(image_a_pred, image_b_pred, non_matches_a, non_matches_b, M=M_descriptor)
  
          non_match_pixel_l2_loss, _, _ = self.l2_pixel_loss(matches_b, non_matches_b, M_pixel=M_pixel)
  
  
  
  
          non_match_loss = (non_match_descriptor_loss * non_match_pixel_l2_loss).sum()
  
          if self.debug:
              self._debug_data['num_hard_negatives'] = num_hard_negatives
              self._debug_data['fraction_hard_negatives'] = num_hard_negatives * 1.0/num_non_matches
  
  
          return non_match_loss, num_hard_negatives
  
      def non_match_loss_descriptor_only(self, image_a_pred, image_b_pred, non_matches_a, non_matches_b, M_descriptor=0.5, invert=False):
          """
          Computes the non-match loss, only using the desciptor norm
          :param image_a_pred:
          :type image_a_pred:
          :param image_b_pred:
          :type image_b_pred:
          :param non_matches_a:
          :type non_matches_a:
          :param non_matches_b:
          :type non_matches_b:
          :param M:
          :type M:
          :return: non_match_loss, num_hard_negatives
          :rtype: torch.Variable, int
          """
          PCL = PixelwiseContrastiveLoss
  
          if M_descriptor is None:
              M_descriptor = self._config["M_descriptor"]
  
          non_match_loss_vec, num_hard_negatives, _, _ = PCL.non_match_descriptor_loss(image_a_pred, image_b_pred, non_matches_a,
                                                                   non_matches_b, M=M_descriptor, invert=invert)
  
          
          num_non_matches = int(non_match_loss_vec.size()[0])
  
  
          non_match_loss = non_match_loss_vec.sum()
  
          if self._debug:
              self._debug_data['num_hard_negatives'] = num_hard_negatives
  
              if num_non_matches == 0:
                  self._debug_data['fraction_hard_negatives'] = 0
              else:
                  self._debug_data['fraction_hard_negatives'] = num_hard_negatives * 1.0/num_non_matches
  
          return non_match_loss, num_hard_negatives
  
  
      def l2_pixel_loss(self, matches_b, non_matches_b, M_pixel=None):
          """
          Apply l2 loss in pixel space.
  
          This weights non-matches more if they are "far away" in pixel space.
  
          :param matches_b: A torch.LongTensor with shape torch.Shape([num_matches])
          :param non_matches_b: A torch.LongTensor with shape torch.Shape([num_non_matches])
          :return l2 loss per sample: A torch.FloatTensorof with shape torch.Shape([num_matches])
          """
  
          if M_pixel is None:
              M_pixel = self._config['M_pixel']
  
          num_non_matches_per_match = len(non_matches_b)/len(matches_b)
  
          ground_truth_pixels_for_non_matches_b = torch.t(matches_b.repeat(num_non_matches_per_match,1)).contiguous().view(-1,1)
  
          ground_truth_u_v_b = self.flattened_pixel_locations_to_u_v(ground_truth_pixels_for_non_matches_b)
          sampled_u_v_b      = self.flattened_pixel_locations_to_u_v(non_matches_b.unsqueeze(1))
  
          # each element is always within [0,1], you have 1 if you are at least M_pixel away in
          # L2 norm in pixel space
          norm_degree = 2
          squared_l2_pixel_loss = 1.0/M_pixel * torch.clamp((ground_truth_u_v_b - sampled_u_v_b).float().norm(norm_degree,1), max=M_pixel)
  
  
          return squared_l2_pixel_loss, ground_truth_u_v_b, sampled_u_v_b
          
  
      
      def flattened_pixel_locations_to_u_v(self, flat_pixel_locations):
          """
          :param flat_pixel_locations: A torch.LongTensor of shape torch.Shape([n,1]) where each element
           is a flattened pixel index, i.e. some integer between 0 and 307,200 for a 640x480 image
  
          :type flat_pixel_locations: torch.LongTensor
  
          :return A torch.LongTensor of shape (n,2) where the first column is the u coordinates of
          the pixel and the second column is the v coordinate
  
          """
          u_v_pixel_locations = flat_pixel_locations.repeat(1,2)
          u_v_pixel_locations[:,0] = u_v_pixel_locations[:,0]%self.image_width 
          u_v_pixel_locations[:,1] = u_v_pixel_locations[:,1]/self.image_width
          return u_v_pixel_locations
  
      def get_l2_pixel_loss_original(self):
          pass
  
      def get_loss_original(self, image_a_pred, image_b_pred, matches_a,
                            matches_b, non_matches_a, non_matches_b,
                            M_margin=0.5, non_match_loss_weight=1.0):
  
          # this is pegged to it's implemenation at sha 87abdb63bb5b99d9632f5c4360b5f6f1cf54245f
          """
          Computes the loss function
          DCN = Dense Correspondence Network
          num_images = number of images in this batch
          num_matches = number of matches
          num_non_matches = number of non-matches
          W = image width
          H = image height
          D = descriptor dimension
          match_loss = 1/num_matches \sum_{num_matches} ||descriptor_a - descriptor_b||_2^2
          non_match_loss = 1/num_non_matches \sum_{num_non_matches} max(0, M_margin - ||descriptor_a - descriptor_b||_2^2 )
          loss = match_loss + non_match_loss
          :param image_a_pred: Output of DCN network on image A.
          :type image_a_pred: torch.Variable(torch.FloatTensor) shape [1, W * H, D]
          :param image_b_pred: same as image_a_pred
          :type image_b_pred:
          :param matches_a: torch.Variable(torch.LongTensor) has shape [num_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of one dimension of image_a_pred
          :type matches_a: torch.Variable(torch.FloatTensor)
          :param matches_b: same as matches_b
          :type matches_b:
          :param non_matches_a: torch.Variable(torch.FloatTensor) has shape [num_non_matches,],  a (u,v) pair is mapped
          to (u,v) ---> image_width * v + u, this matches the shape of image_a_pred
          :type non_matches_a: torch.Variable(torch.FloatTensor)
          :param non_matches_b: same as non_matches_a
          :type non_matches_b:
          :return: loss, match_loss, non_match_loss
          :rtype: torch.Variable(torch.FloatTensor) each of shape torch.Size([1])
          """
  
          num_matches = matches_a.size()[0]
          num_non_matches = non_matches_a.size()[0]
  
  
          matches_a_descriptors = torch.index_select(image_a_pred, 1, matches_a)
          matches_b_descriptors = torch.index_select(image_b_pred, 1, matches_b)
  
          match_loss = 1.0/num_matches * (matches_a_descriptors - matches_b_descriptors).pow(2).sum()
  
          # add loss via non_matches
          non_matches_a_descriptors = torch.index_select(image_a_pred, 1, non_matches_a)
          non_matches_b_descriptors = torch.index_select(image_b_pred, 1, non_matches_b)
          pixel_wise_loss = (non_matches_a_descriptors - non_matches_b_descriptors).pow(2).sum(dim=2)
          pixel_wise_loss = torch.add(torch.neg(pixel_wise_loss), M_margin)
          zeros_vec = torch.zeros_like(pixel_wise_loss)
          non_match_loss = non_match_loss_weight * 1.0/num_non_matches * torch.max(zeros_vec, pixel_wise_loss).sum()
  
          loss = match_loss + non_match_loss
  
          return loss, match_loss, non_match_loss
  