adalam.py

import torch

from ..utils.base_model import BaseModel

from kornia.feature.adalam import AdalamFilter
from kornia.utils.helpers import get_cuda_device_if_available


class AdaLAM(BaseModel):
    # See https://kornia.readthedocs.io/en/latest/_modules/kornia/feature/adalam/adalam.html.
    default_conf = {
        'area_ratio': 100,
        'search_expansion': 4,
        'ransac_iters': 128,
        'min_inliers': 6,
        'min_confidence': 200,
        'orientation_difference_threshold': 30,
        'scale_rate_threshold': 1.5,
        'detected_scale_rate_threshold': 5,
        'refit': True,
        'force_seed_mnn': True,
        'device': get_cuda_device_if_available()
    }
    required_inputs = [
        'image0', 'image1',
        'descriptors0', 'descriptors1',
        'keypoints0', 'keypoints1',
        'scales0', 'scales1',
        'oris0', 'oris1']

    def _init(self, conf):
        self.adalam = AdalamFilter(conf)

    def _forward(self, data):
        assert data['keypoints0'].size(0) == 1
        if data['keypoints0'].size(1) < 2 or data['keypoints1'].size(1) < 2:
            matches = torch.zeros(
                (0, 2), dtype=torch.int64,
                device=data['keypoints0'].device)
        else:
            matches = self.adalam.match_and_filter(
                data['keypoints0'][0], data['keypoints1'][0],
                data['descriptors0'][0].T, data['descriptors1'][0].T,
                data['image0'].shape[2:], data['image1'].shape[2:],
                data['oris0'][0], data['oris1'][0],
                data['scales0'][0], data['scales1'][0]
            )
        matches_new = torch.full(
            (data['keypoints0'].size(1),), -1,
            dtype=torch.int64, device=data['keypoints0'].device)
        matches_new[matches[:, 0]] = matches[:, 1]
        return {
            'matches0': matches_new.unsqueeze(0),
            'matching_scores0': torch.zeros(matches_new.size(0)).unsqueeze(0)
        }