docstr: Args:
      conf_matrix (torch.Tensor): [N, L, S]
      data (dict): with keys ['hw0_i', 'hw1_i', 'hw0_c', 'hw1_c']
  Returns:
      coarse_matches (dict): {
          'b_ids' (torch.Tensor): [M'],
          'i_ids' (torch.Tensor): [M'],
          'j_ids' (torch.Tensor): [M'],
          'gt_mask' (torch.Tensor): [M'],
          'm_bids' (torch.Tensor): [M],
          'mkpts0_c' (torch.Tensor): [M, 2],
          'mkpts1_c' (torch.Tensor): [M, 2],
          'mconf' (torch.Tensor): [M]}
docstr_tok: ["Args",":","conf_matrix","(","torch.Tensor",")",":","[","N",",","L",",","S","]","data","(","dict",")",":","with","keys","[","'hw0_i","'",",","'hw1_i","'",",","'hw0_c","'",",","'hw1_c","'","]","Returns",":","coarse_matches","(","dict",")",":","{","'b_ids","'","(","torch.Tensor",")",":","[","M","'","]",",","'i_ids","'","(","torch.Tensor",")",":","[","M","'","]",",","'j_ids","'","(","torch.Tensor",")",":","[","M","'","]",",","'gt_mask","'","(","torch.Tensor",")",":","[","M","'","]",",","'m_bids","'","(","torch.Tensor",")",":","[","M","]",",","'mkpts0_c","'","(","torch.Tensor",")",":","[","M",",","2","]",",","'mkpts1_c","'","(","torch.Tensor",")",":","[","M",",","2","]",",","'mconf","'","(","torch.Tensor",")",":","[","M","]","}"]
code: def get_coarse_match(
      self, conf_matrix: Tensor, data: Dict[str, Tensor]
  ) -> Dict[str, Tensor]:
      """
      Args:
          conf_matrix (torch.Tensor): [N, L, S]
          data (dict): with keys ['hw0_i', 'hw1_i', 'hw0_c', 'hw1_c']
      Returns:
          coarse_matches (dict): {
              'b_ids' (torch.Tensor): [M'],
              'i_ids' (torch.Tensor): [M'],
              'j_ids' (torch.Tensor): [M'],
              'gt_mask' (torch.Tensor): [M'],
              'm_bids' (torch.Tensor): [M],
              'mkpts0_c' (torch.Tensor): [M, 2],
              'mkpts1_c' (torch.Tensor): [M, 2],
              'mconf' (torch.Tensor): [M]}
      """
      axes_lengths = {
          "h0c": data["hw0_c"][0],
          "w0c": data["hw0_c"][1],
          "h1c": data["hw1_c"][0],
          "w1c": data["hw1_c"][1],
      }
      _device = conf_matrix.device
      # 1. confidence thresholding
      mask = conf_matrix > self.thr
      N = conf_matrix.shape[0]
      mask = mask.reshape(
          N,
          axes_lengths["h0c"],
          axes_lengths["w0c"],
          axes_lengths["h1c"],
          axes_lengths["w1c"],
      )
      # mask = rearrange(mask, 'b (h0c w0c) (h1c w1c) -> b h0c w0c h1c w1c',
      #                 **axes_lengths)
      if "mask0" not in data:
          mask_border(mask, self.border_rm, False)
      else:
          mask_border_with_padding(
              mask, self.border_rm, False, data["mask0"], data["mask1"]
          )
      mask = mask.reshape(
          N,
          axes_lengths["h0c"] * axes_lengths["w0c"],
          axes_lengths["h1c"] * axes_lengths["w1c"],
      )
      # rearrange(mask, 'b h0c w0c h1c w1c -> b (h0c w0c) (h1c w1c)',
      #                 **axes_lengths)
  
      # 2. mutual nearest
      mask = (
          mask
          * (conf_matrix == conf_matrix.max(dim=2, keepdim=True)[0])
          * (conf_matrix == conf_matrix.max(dim=1, keepdim=True)[0])
      )
  
      # 3. find all valid coarse matches
      # this only works when at most one `True` in each row
      mask_v, all_j_ids = mask.max(dim=2)
      b_ids, i_ids = torch.where(mask_v)
      j_ids = all_j_ids[b_ids, i_ids]
      mconf = conf_matrix[b_ids, i_ids, j_ids]
  
      # 4. Random sampling of training samples for fine-level LoFTR
      # (optional) pad samples with gt coarse-level matches
      if self.training:
          # NOTE:
          # The sampling is performed across all pairs in a batch without manually balancing
          # #samples for fine-level increases w.r.t. batch_size
          if "mask0" not in data:
              num_candidates_max = mask.size(0) * max(
                  mask.size(1), mask.size(2)
              )
          else:
              num_candidates_max = compute_max_candidates(
                  data["mask0"], data["mask1"]
              )
          num_matches_train = (
              num_candidates_max * self.train_coarse_percent
          )
          num_matches_train = int(num_matches_train)
          num_matches_pred = len(b_ids)
          if self.train_pad_num_gt_min >= num_matches_train:
              msg = (
                  "min-num-gt-pad should be less than num-train-matches"
              )
              raise ValueError(msg)
  
          # pred_indices is to select from prediction
          if (
              num_matches_pred
              <= num_matches_train - self.train_pad_num_gt_min
          ):
              pred_indices = torch.arange(
                  num_matches_pred, device=_device
              )
          else:
              pred_indices = torch.randint(
                  num_matches_pred,
                  (num_matches_train - self.train_pad_num_gt_min,),
                  device=_device,
              )
  
          # gt_pad_indices is to select from gt padding. e.g. max(3787-4800, 200)
          gt_pad_indices = torch.randint(
              len(data["spv_b_ids"]),
              (
                  max(
                      num_matches_train - num_matches_pred,
                      self.train_pad_num_gt_min,
                  ),
              ),
              device=_device,
          )
          mconf_gt = torch.zeros(
              len(data["spv_b_ids"]), device=_device
          )  # set conf of gt paddings to all zero
  
          b_ids, i_ids, j_ids, mconf = (  # type: ignore
              torch.cat([x[pred_indices], y[gt_pad_indices]], dim=0)
              for x, y in zip(
                  [b_ids, data["spv_b_ids"]],
                  [i_ids, data["spv_i_ids"]],
                  [j_ids, data["spv_j_ids"]],
                  [mconf, mconf_gt],
              )
          )
  
      # These matches select patches that feed into fine-level network
      coarse_matches = {"b_ids": b_ids, "i_ids": i_ids, "j_ids": j_ids}
  
      # 4. Update with matches in original image resolution
      scale = data["hw0_i"][0] / data["hw0_c"][0]
      scale0 = (
          scale * data["scale0"][b_ids] if "scale0" in data else scale
      )
      scale1 = (
          scale * data["scale1"][b_ids] if "scale1" in data else scale
      )
      mkpts0_c = (
          torch.stack(
              [i_ids % data["hw0_c"][1], i_ids // data["hw0_c"][1]],
              dim=1,
          )
          * scale0
      )
      mkpts1_c = (
          torch.stack(
              [j_ids % data["hw1_c"][1], j_ids // data["hw1_c"][1]],
              dim=1,
          )
          * scale1
      )
  
      # These matches is the current prediction (for visualization)
      coarse_matches.update(
          {
              "gt_mask": mconf == 0,
              "m_bids": b_ids[mconf != 0],  # mconf == 0 => gt matches
              "mkpts0_c": mkpts0_c[mconf != 0].to(
                  dtype=conf_matrix.dtype
              ),
              "mkpts1_c": mkpts1_c[mconf != 0].to(
                  dtype=conf_matrix.dtype
              ),
              "mconf": mconf[mconf != 0],
          }
      )
  
      return coarse_matches
  
file_path: kornia/kornia/feature/loftr/utils/coarse_matching.py
imports: from typing import Any, Dict, Optional, Union
  import torch
  import torch.nn.functional
  from torch import nn
  from kornia.core import Module, Tensor
label_desc: Use image node_id 3 for calling the CoarseMatching obj's underlying member method code with example usage: obj.get_coarse_match(conf_matrix, data) and returns: coarse_matches