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]
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

    
kornia/feature/adalam/core.py