nearest_neighbor.py

import torch

from ..utils.base_model import BaseModel


def find_nn(sim, ratio_thresh, distance_thresh):
    sim_nn, ind_nn = sim.topk(2 if ratio_thresh else 1, dim=-1, largest=True)
    dist_nn = 2 * (1 - sim_nn)
    mask = torch.ones(ind_nn.shape[:-1], dtype=torch.bool, device=sim.device)
    if ratio_thresh:
        mask = mask & (dist_nn[..., 0] <= (ratio_thresh**2)*dist_nn[..., 1])
    if distance_thresh:
        mask = mask & (dist_nn[..., 0] <= distance_thresh**2)
    matches = torch.where(mask, ind_nn[..., 0], ind_nn.new_tensor(-1))
    scores = torch.where(mask, (sim_nn[..., 0]+1)/2, sim_nn.new_tensor(0))
    return matches, scores


def mutual_check(m0, m1):
    inds0 = torch.arange(m0.shape[-1], device=m0.device)
    loop = torch.gather(m1, -1, torch.where(m0 > -1, m0, m0.new_tensor(0)))
    ok = (m0 > -1) & (inds0 == loop)
    m0_new = torch.where(ok, m0, m0.new_tensor(-1))
    return m0_new


class NearestNeighbor(BaseModel):
    default_conf = {
        'ratio_threshold': None,
        'distance_threshold': None,
        'do_mutual_check': True,
    }
    required_inputs = ['descriptors0', 'descriptors1']

    def _init(self, conf):
        pass

    def _forward(self, data):
        if data['descriptors0'].size(-1) == 0 or data['descriptors1'].size(-1) == 0:
            matches0 = torch.full(
                data['descriptors0'].shape[:2], -1,
                device=data['descriptors0'].device)
            return {
                'matches0': matches0,
                'matching_scores0': torch.zeros_like(matches0)
            }
        ratio_threshold = self.conf['ratio_threshold']
        if data['descriptors0'].size(-1) == 1 or data['descriptors1'].size(-1) == 1:
            ratio_threshold = None
        sim = torch.einsum(
            'bdn,bdm->bnm', data['descriptors0'], data['descriptors1'])
        matches0, scores0 = find_nn(
            sim, ratio_threshold, self.conf['distance_threshold'])
        if self.conf['do_mutual_check']:
            matches1, scores1 = find_nn(
                sim.transpose(1, 2), ratio_threshold,
                self.conf['distance_threshold'])
            matches0 = mutual_check(matches0, matches1)
        return {
            'matches0': matches0,
            'matching_scores0': scores0,
        }