text: # -*- coding: utf-8 -*-
  # Copyright (c) Facebook, Inc. and its affiliates.
  # Modified by BaseDetection, Inc. and its affiliates.
  from scipy.optimize import linear_sum_assignment
  
  import torch
  from torch import nn
  
  from cvpods.structures.boxes import box_cxcywh_to_xyxy, generalized_box_iou
  
  
  class Matcher(object):
      """
      This class assigns to each predicted "element" (e.g., a box) a ground-truth
      element. Each predicted element will have exactly zero or one matches; each
      ground-truth element may be matched to zero or more predicted elements.
  
      The matching is determined by the MxN match_quality_matrix, that characterizes
      how well each (ground-truth, prediction)-pair match each other. For example,
      if the elements are boxes, this matrix may contain box intersection-over-union
      overlap values.
  
      The matcher returns (a) a vector of length N containing the index of the
      ground-truth element m in [0, M) that matches to prediction n in [0, N).
      (b) a vector of length N containing the labels for each prediction.
      """
  
      def __init__(self, thresholds, labels, allow_low_quality_matches=False):
          """
          Args:
              thresholds (list): a list of thresholds used to stratify predictions
                  into levels.
              labels (list): a list of values to label predictions belonging at
                  each level. A label can be one of {-1, 0, 1} signifying
                  {ignore, negative class, positive class}, respectively.
              allow_low_quality_matches (bool): if True, produce additional matches
                  for predictions with maximum match quality lower than high_threshold.
                  See set_low_quality_matches_ for more details.
  
              For example,
                  thresholds = [0.3, 0.5]
                  labels = [0, -1, 1]
                  All predictions with iou < 0.3 will be marked with 0 and
                  thus will be considered as false positives while training.
                  All predictions with 0.3 <= iou < 0.5 will be marked with -1 and
                  thus will be ignored.
                  All predictions with 0.5 <= iou will be marked with 1 and
                  thus will be considered as true positives.
          """
          # Add -inf and +inf to first and last position in thresholds
          thresholds = thresholds[:]
          assert thresholds[0] > 0
          thresholds.insert(0, -float("inf"))
          thresholds.append(float("inf"))
          assert all(low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:]))
          assert all(label in [-1, 0, 1] for label in labels)
          assert len(labels) == len(thresholds) - 1
          self.thresholds = thresholds
          self.labels = labels
          self.allow_low_quality_matches = allow_low_quality_matches
  
      def __call__(self, match_quality_matrix):
          """
          Args:
              match_quality_matrix (Tensor[float]): an MxN tensor, containing the
                  pairwise quality between M ground-truth elements and N predicted
                  elements. All elements must be >= 0 (due to the us of `torch.nonzero`
                  for selecting indices in :meth:`set_low_quality_matches_`).
  
          Returns:
              matches (Tensor[int64]): a vector of length N, where matches[i] is a matched
                  ground-truth index in [0, M)
              match_labels (Tensor[int8]): a vector of length N, where pred_labels[i] indicates
                  whether a prediction is a true or false positive or ignored
          """
          assert match_quality_matrix.dim() == 2
          if match_quality_matrix.numel() == 0:
              default_matches = match_quality_matrix.new_full(
                  (match_quality_matrix.size(1),), 0, dtype=torch.int64
              )
              # When no gt boxes exist, we define IOU = 0 and therefore set labels
              # to `self.labels[0]`, which usually defaults to background class 0
              # To choose to ignore instead, can make labels=[-1,0,-1,1] + set appropriate thresholds
              default_match_labels = match_quality_matrix.new_full(
                  (match_quality_matrix.size(1),), self.labels[0], dtype=torch.int8
              )
              return default_matches, default_match_labels
  
          assert torch.all(match_quality_matrix >= 0)
  
          # match_quality_matrix is M (gt) x N (predicted)
          # Max over gt elements (dim 0) to find best gt candidate for each prediction
          matched_vals, matches = match_quality_matrix.max(dim=0)
  
          match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)
  
          for (l, low, high) in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):
              low_high = (matched_vals >= low) & (matched_vals < high)
              match_labels[low_high] = l
  
          if self.allow_low_quality_matches:
              self.set_low_quality_matches_(match_labels, match_quality_matrix)
  
          return matches, match_labels
  
      def set_low_quality_matches_(self, match_labels, match_quality_matrix):
          """
          Produce additional matches for predictions that have only low-quality matches.
          Specifically, for each ground-truth G find the set of predictions that have
          maximum overlap with it (including ties); for each prediction in that set, if
          it is unmatched, then match it to the ground-truth G.
  
          This function implements the RPN assignment case (i) in Sec. 3.1.2 of the
          Faster R-CNN paper: https://arxiv.org/pdf/1506.01497v3.pdf.
          """
          # For each gt, find the prediction with which it has highest quality
          highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
          # Find the highest quality match available, even if it is low, including ties.
          # Note that the matches qualities must be positive due to the use of
          # `torch.nonzero`.
          gt_pred_pairs_of_highest_quality = torch.nonzero(
              match_quality_matrix == highest_quality_foreach_gt[:, None],
              as_tuple=False
          )
          # Example gt_pred_pairs_of_highest_quality:
          #   tensor([[    0, 39796],
          #           [    1, 32055],
          #           [    1, 32070],
          #           [    2, 39190],
          #           [    2, 40255],
          #           [    3, 40390],
          #           [    3, 41455],
          #           [    4, 45470],
          #           [    5, 45325],
          #           [    5, 46390]])
          # Each row is a (gt index, prediction index)
          # Note how gt items 1, 2, 3, and 5 each have two ties
  
          pred_inds_to_update = gt_pred_pairs_of_highest_quality[:, 1]
          match_labels[pred_inds_to_update] = 1
  
  
  class MatcherIgnore(Matcher):
      """
      Matcher with Ignore labels. (e.g. Crowd Human)
      This class assigns to each predicted "element" (e.g., a box) a ground-truth
      element. Each predicted element will have exactly zero or one matches; each
      ground-truth element may be matched to zero or more predicted elements.
  
      The matching is determined by the MxN match_quality_matrix, that characterizes
      how well each (ground-truth, prediction)-pair match each other. For example,
      if the elements are boxes, this matrix may contain box intersection-over-union
      overlap values.
  
      The matcher returns (a) a vector of length N containing the index of the
      ground-truth element m in [0, M) that matches to prediction n in [0, N).
      (b) a vector of length N containing the labels for each prediction.
      """
  
      def __call__(self, match_quality_matrix, match_quality_ignore=None):
          """
          Args:
              match_quality_matrix (Tensor[float]): an MxN tensor, containing the
                  pairwise quality between M ground-truth elements and N predicted
                  elements. All elements must be >= 0 (due to the us of `torch.nonzero`
                  for selecting indices in :meth:`set_low_quality_matches_`).
  
          Returns:
              matches (Tensor[int64]): a vector of length N, where matches[i] is a matched
                  ground-truth index in [0, M)
              match_labels (Tensor[int8]): a vector of length N, where pred_labels[i] indicates
                  whether a prediction is a true or false positive or ignored
          """
          assert match_quality_matrix.dim() == 2
          if match_quality_matrix.numel() == 0:
              default_matches = match_quality_matrix.new_full(
                  (match_quality_matrix.size(1),), 0, dtype=torch.int64
              )
              # When no gt boxes exist, we define IOU = 0 and therefore set labels
              # to `self.labels[0]`, which usually defaults to background class 0
              # To choose to ignore instead, can make labels=[-1,0,-1,1] + set appropriate thresholds
              default_match_labels = match_quality_matrix.new_full(
                  (match_quality_matrix.size(1),), self.labels[0], dtype=torch.int8
              )
              return default_matches, default_match_labels
  
          assert torch.all(match_quality_matrix >= 0)
  
          # match_quality_matrix is M (gt) x N (predicted)
          # Max over gt elements (dim 0) to find best gt candidate for each prediction
          matched_vals, matches = match_quality_matrix.max(dim=0)
          matched_vals_ign, matches_ign = match_quality_ignore.max(dim=0)
  
          assert len(self.thresholds) == 3
          high = self.thresholds[1]
          ign_mask = ((matched_vals_ign > matched_vals) & (matched_vals < high)).float()
          matched_vals = matched_vals * (1 - ign_mask) + matched_vals_ign * ign_mask
          matches = (matches * (1 - ign_mask) + matches_ign * ign_mask).long()
  
          match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)
  
          for (l, low, high) in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):
              low_high = (matched_vals >= low) & (matched_vals < high)
              match_labels[low_high] = l
  
          if self.allow_low_quality_matches:
              self.set_low_quality_matches_(match_labels, match_quality_matrix)
  
          return matches, match_labels
  
  
  class HungarianMatcher(nn.Module):
      """This class computes an assignment between the targets and the predictions of the network
      For efficiency reasons, the targets don't include the no_object. Because of this, in general,
      there are more predictions than targets. In this case, we do a 1-to-1 matching of the best
      predictions, while the others are un-matched (and thus treated as non-objects).
      """
  
      def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1):
          """Creates the matcher
          Params:
              cost_class: This is the relative weight of the classification error in the matching cost
              cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates
              in the matching cost cost_giou: This is the relative weight of the giou loss of the
              bounding box in the matching cost
          """
          super().__init__()
          self.cost_class = cost_class
          self.cost_bbox = cost_bbox
          self.cost_giou = cost_giou
          assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, "all costs cant be 0"
  
      @torch.no_grad()
      def forward(self, outputs, targets):
          """ Performs the matching
          Params:
              outputs: This is a dict that contains at least these entries:
                   "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the
                              classification logits
                   "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted
                              box coordinates
              targets: This is a list of targets (len(targets) = batch_size), where each target
                              is a dict containing:
                   "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number
                              of ground-truth objects in the target) containing the class labels
                   "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates
          Returns:
              A list of size batch_size, containing tuples of (index_i, index_j) where:
                  - index_i is the indices of the selected predictions (in order)
                  - index_j is the indices of the corresponding selected targets (in order)
              For each batch element, it holds:
                  len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
          """
          bs, num_queries = outputs["pred_logits"].shape[:2]
  
          # We flatten to compute the cost matrices in a batch
          out_prob = (
              outputs["pred_logits"].flatten(0, 1).softmax(-1)
          )  # [batch_size * num_queries, num_classes]
          out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]
  
          # Also concat the target labels and boxes
          tgt_ids = torch.cat([v["labels"] for v in targets])
          tgt_bbox = torch.cat([v["boxes"] for v in targets])
  
          # Compute the classification cost. Contrary to the loss, we don't use the NLL,
          # but approximate it in 1 - proba[target class].
          # The 1 is a constant that doesn't change the matching, it can be ommitted.
          cost_class = -out_prob[:, tgt_ids]
  
          # Compute the L1 cost between boxes
          cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)
  
          # Compute the giou cost betwen boxes
          cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))
  
          # Final cost matrix
          C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou
          C = C.view(bs, num_queries, -1).cpu()
  
          sizes = [len(v["boxes"]) for v in targets]
          indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
          return [
              (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64))
              for i, j in indices
          ]
  