alike_wrapper.py

"""
	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
"""


import sys
import os

ALIKE_PATH = (os.path.abspath(os.path.dirname(__file__) + '/ALIKE'))
sys.path.append(ALIKE_PATH)

import torch
from alike import ALike
import cv2
import numpy as np

import pdb

dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

configs = {
    'alike-t': {'c1': 8, 'c2': 16, 'c3': 32, 'c4': 64, 'dim': 64, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-t.pth')},
    'alike-s': {'c1': 8, 'c2': 16, 'c3': 48, 'c4': 96, 'dim': 96, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-s.pth')},
    'alike-n': {'c1': 16, 'c2': 32, 'c3': 64, 'c4': 128, 'dim': 128, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-n.pth')},
    'alike-l': {'c1': 32, 'c2': 64, 'c3': 128, 'c4': 128, 'dim': 128, 'single_head': False, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-l.pth')},
}

model = ALike(**configs['alike-t'],
                device=dev,
                top_k=4096,
                scores_th=0.1,
                n_limit=8000)

def extract_alike_kpts(img):
    pred0 = model(img, sub_pixel=True)
    return pred0['keypoints']

def detectAndCompute(img, top_k = 4096):

    img = (img[0].permute(1,2,0).cpu().numpy() * 255).astype(np.uint8)

    pred0 = model(img, sub_pixel=True)
    return (torch.tensor(pred0['keypoints'], dtype=torch.float32), 
            torch.tensor(pred0['scores'], dtype=torch.float32), 
            torch.tensor(pred0['descriptors'], dtype = torch.float32)
            )

def match_alike(img1, img2):
    from kornia.feature import DescriptorMatcher
    kornia_matcher = DescriptorMatcher('mnn')
    with torch.inference_mode():
        pred0 = model(img1, sub_pixel=True)
        pred1 = model(img2, sub_pixel=True)

        kp1 = pred0['keypoints']
        kp2 = pred1['keypoints']
        des1 = pred0['descriptors']
        des2 = pred1['descriptors']

        des1 = torch.tensor(des1, device=dev)
        des2 = torch.tensor(des2, device=dev)
        dists, matches = kornia_matcher(des1, des2)

        mkpts0 = kp1[matches[:,0].cpu().numpy()]
        mkpts1 = kp2[matches[:,1].cpu().numpy()]

    return mkpts0, mkpts1
 

def create_xy(h, w, dev):
    y, x = torch.meshgrid(torch.arange(h, device = dev), 
                            torch.arange(w, device = dev))
    xy = torch.cat([x[..., None],y[..., None]], -1).reshape(-1,2)
    return xy

def match_alike_customkp(img1, img2, kp_img1):
    from kornia.feature import DescriptorMatcher
    kornia_matcher = DescriptorMatcher('nn')
    with torch.inference_mode():
        kp_img1 = kp_img1.astype(np.int32)
        pred0 = model(img1, sub_pixel=True, return_dense=True)
        pred1 = model(img2, sub_pixel=True, return_dense=True)

        print(pred0.keys())

        des1 = pred0['desc_map'].cpu()
        des2 = pred1['desc_map'].cpu()

        B, C, H, W = des1.shape
        kp2  = create_xy(H, W, des1.device).numpy()
        des2 = des2.reshape(C, -1).permute(1,0)
        des1 = des1[0, :, kp_img1[:,1], kp_img1[:,0]].permute(1,0)
        kp1 = kp_img1

        print(des1.shape, des2.shape, kp1.shape, kp2.shape)
        dists, matches = kornia_matcher(des1, des2)
        #print('time: ', pred1['time'])


    # Extract matching keypoints' positions
    mkpts0 = kp1[matches[:,0]]
    mkpts1 = kp2[matches[:,1]]

    return mkpts0.astype(np.float32), mkpts1.astype(np.float32)