code: def extract_neighborhood_sets(
      o1: Optional[Tensor],
      o2: Optional[Tensor],
      s1: Optional[Tensor],
      s2: Optional[Tensor],
      dist1: Tensor,
      im1seeds: Tensor,
      im2seeds: Tensor,
      k1: Tensor,
      k2: Tensor,
      R1: Union[float, Tensor],
      R2: Union[float, Tensor],
      fnn12: Tensor,
      ORIENTATION_THR: float,
      SCALE_RATE_THR: float,
      SEARCH_EXP: float,
      MIN_INLIERS: float,
  ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
      """Assign keypoints to seed points. This checks both the distance and the agreement of the local transformation
      if available.
  
      o1: Orientations of keypoints in image I_1
      o2: Orientations of keypoints in image I_2
      s1: Scales of keypoints in image I_1
      s2: Scales of keypoints in image I_2
      dist1: Precomputed distance matrix between keypoints in image I_1
      im1seeds: Keypoint index of chosen seeds in image I_1
      im2seeds: Keypoint index of chosen seeds in image I_2
      k1: Keypoint locations in image I_1
      k2: Keypoint locations in image I_2
      R1: Base radius of neighborhoods in image I_1
      R2: Base radius of neighborhoods in image I_2
      fnn12: Matches between keypoints of I_1 and I_2.
             The i-th entry of fnn12 is j if and only if keypoint k_i in image I_1 is matched to keypoint k_j in image I_2
      ORIENTATION_THR: Maximum deviation of orientation with respect to seed S_i to keep a keypoint in i-th neighborhood
      SCALE_RATE_THR: Maximum deviation of scale with respect to seed S_i to keep a keypoint in i-th neighborhood
      SEARCH_EXP: Expansion rate for both radii R1 and R2 to consider inclusion of neighboring keypoints
      MIN_INLIERS: Minimum number of inliers to keep a seed point. This is used as an early filter here
                   to remove already seeds with not enough samples to ever pass this threshold.
  
      Returns:
          Local neighborhoods assignments:
  
          local_neighs_mask: Boolean matrix of size (num_seeds, num_keypoints).
                             Entry (i, j) is True iff keypoint j was assigned to seed i.
          rdims: Number of keypoints included in the neighborhood for each seed
          im1seeds: Keypoint index of chosen seeds in image I_1
          im2seeds: Keypoint index of chosen seeds in image I_2
      """
      dst1 = dist1[im1seeds, :]
      dst2 = dist_matrix(k2[fnn12[im1seeds]], k2[fnn12])
  
      # initial candidates are matches which are close to the same seed in both images
      local_neighs_mask = (dst1 < (SEARCH_EXP * R1) ** 2) & (
          dst2 < (SEARCH_EXP * R2) ** 2
      )
  
      # If requested, also their orientation delta should be compatible with that of the corresponding seed
      if (
          ORIENTATION_THR is not None
          and ORIENTATION_THR < 180
          and (o1 is not None)
          and (o2 is not None)
      ):
          relo = orientation_diff(o1, o2[fnn12])
          orientation_diffs = torch.abs(
              orientation_diff(
                  relo.unsqueeze(0), relo[im1seeds].unsqueeze(1)
              )
          )
          local_neighs_mask = local_neighs_mask & (
              orientation_diffs < ORIENTATION_THR
          )
  
      # If requested, also their scale delta should be compatible with that of the corresponding seed
      if (
          SCALE_RATE_THR is not None
          and (SCALE_RATE_THR < 10)
          and (s1 is not None)
          and (s2 is not None)
      ):
          rels = s2[fnn12] / s1
          scale_rates = rels[im1seeds].unsqueeze(1) / rels.unsqueeze(0)
          local_neighs_mask = (
              local_neighs_mask
              & (scale_rates < SCALE_RATE_THR)
              & (scale_rates > 1 / SCALE_RATE_THR)
          )  # (ns, n1)
  
      # count how many keypoints ended up in each neighborhood
      numn1 = torch.sum(local_neighs_mask, dim=1)
      # and only keep the ones that have enough points
      valid_seeds = numn1 >= MIN_INLIERS
  
      local_neighs_mask = local_neighs_mask[valid_seeds, :]
  
      rdims = numn1[valid_seeds]
  
      return (
          local_neighs_mask,
          rdims,
          im1seeds[valid_seeds],
          im2seeds[valid_seeds],
      )
  
file_path: kornia/kornia/feature/adalam/core.py
imports: import math
  from typing import Optional, Tuple, Union
  import torch
  from typing_extensions import NotRequired, TypedDict
  from kornia.core import Tensor, concatenate, tensor, where
  from .ransac import ransac
  from .utils import dist_matrix, orientation_diff
label_desc: Use image node_id 3 for calling a global function with example usage: extract_neighborhood_sets(o1, o2, s1, s2, dist1, im1seeds, im2seeds, k1, k2, R1, R2, fnn12, ORIENTATION_THR, SCALE_RATE_THR, SEARCH_EXP, MIN_INLIERS) and returns: local_neighs_mask,  rdims,  im1seeds,  im2seeds