vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/keypt2subpx/dataprocess/xfeat.py

@@ -0,0 +1,187 @@
+import argparse
+import torch
+from torch.nn import functional as F
+import numpy as np
+import os
+import pickle
+from tqdm import tqdm
+
+from omegaconf import OmegaConf
+
+from gluefactory.datasets import get_dataset
+from gluefactory.utils.tensor import batch_to_device
+from gluefactory.utils.patches import extract_patches
+
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.enabled = True
+torch.backends.cudnn.benchmark = True
+
+# parse command line arguments
+parser = argparse.ArgumentParser(
+	description='Preprocess MegaDepth dataset with XFeat.',
+	formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+parser.add_argument('--variant', '-v', default='train', choices=['train', 'val', 'test'],
+	help='Defines subfolders of the dataset ot use (split according to "Glue Factory", https://github.com/cvg/glue-factory).')
+
+parser.add_argument('--topk', '-tk', type=int, default=4096, 
+	help='number of features per image')
+
+parser.add_argument('--patch_radius', '-ps', type=int, default=5, 
+	help='patch radius')
+
+parser.add_argument('--dataconf', '-dc', type=str, default='./configs/megadepth_dataconf_b1_grey.yaml',
+	help='data configuration file path')
+
+opt = parser.parse_args()
+
+# setup dataset & dataloader
+dataconf = OmegaConf.load(opt.dataconf)
+dataset = get_dataset(dataconf.name)(dataconf)
+dataloader = dataset.get_data_loader(opt.variant, distributed=False, shuffle=False)
+print(f'Using dataset: {opt.dataconf} {opt.variant}', flush=True)
+
+# output folder that stores pre-calculated correspondence vectors as PyTorch tensors
+out_dir = 'data/megadepth/' + opt.variant + '_xfeat/'
+if not os.path.isdir(out_dir): os.makedirs(out_dir)
+
+# setup detector
+xfeat = torch.hub.load('verlab/accelerated_features', 'XFeat', pretrained = True, top_k = opt.topk).cuda()
+xfeat.eval()
+
+# The matcher part is borrowed and modified from GlueFactory (https://github.com/cvg/glue-factory)
+# Nearest neighbor matcher for normalized descriptors.
+# Optionally apply the mutual check and threshold the distance or ratio.
+@torch.no_grad()
+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))
+    return matches
+
+
+def mutual_check(m0, m1):
+    inds0 = torch.arange(m0.shape[-1], device=m0.device)
+    inds1 = torch.arange(m1.shape[-1], device=m1.device)
+    loop0 = torch.gather(m1, -1, torch.where(m0 > -1, m0, m0.new_tensor(0)))
+    loop1 = torch.gather(m0, -1, torch.where(m1 > -1, m1, m1.new_tensor(0)))
+    m0_new = torch.where((m0 > -1) & (inds0 == loop0), m0, m0.new_tensor(-1))
+    m1_new = torch.where((m1 > -1) & (inds1 == loop1), m1, m1.new_tensor(-1))
+    return m0_new, m1_new
+
+ratio_thresh = None
+distance_thresh = None
+
+def nearest_neighbor_match(descriptors0, descriptors1):
+    sim = torch.einsum("bnd,bmd->bnm", descriptors0, descriptors1)
+    matches0 = find_nn(sim, ratio_thresh, distance_thresh)
+    matches1 = find_nn(
+        sim.transpose(1, 2), ratio_thresh, distance_thresh
+    )
+    matches0, matches1 = mutual_check(matches0, matches1)
+    b, m, n = sim.shape
+    la = sim.new_zeros(b, m + 1, n + 1)
+    la[:, :-1, :-1] = F.log_softmax(sim, -1) + F.log_softmax(sim, -2)
+    mscores0 = (matches0 > -1).float()
+    mscores1 = (matches1 > -1).float()
+    return {
+        "matches0": matches0,
+        "matches1": matches1,
+        "matching_scores0": mscores0,
+        "matching_scores1": mscores1,
+        "similarity": sim,
+        "log_assignment": la,
+    }
+# Borrowed part ends
+
+with torch.no_grad():
+    for i, data in enumerate(tqdm(dataloader)):
+        scene = data['view0']['scene'][0]
+        assert data['view0']['scene'][0] == data['view1']['scene'][0], "Scenes' names are not the same"
+        img1_name = data['view0']['name'][0].split('.')[0]
+        img2_name = data['view1']['name'][0].split('.')[0]
+
+        img1_shape = data['view0']['image_size'][0].numpy() # W, H
+        img1_shape = np.array([img1_shape[1], img1_shape[0], 3]) # H, W, 3
+        img2_shape = data['view1']['image_size'][0].numpy() # W, H
+        img2_shape = np.array([img2_shape[1], img2_shape[0], 3]) # H, W, 3
+
+        tqdm.write("Processing pair %5d of %5d. (%30s, %30s)" % (i, len(dataloader), img1_name, img2_name))
+
+        gpu_data = batch_to_device(data, 'cuda:0', non_blocking=True)
+
+        output = xfeat.detectAndCompute(gpu_data['view0']['image'][:,:,:img1_shape[0],:img1_shape[1]], top_k=4096)[0]
+        pts1, score1, desc1 = output['keypoints'], output['scores'], output['descriptors']
+        output = xfeat.detectAndCompute(gpu_data['view1']['image'][:,:,:img2_shape[0],:img2_shape[1]], top_k=4096)[0]
+        pts2, score2, desc2 = output['keypoints'], output['scores'], output['descriptors']
+
+        pred = nearest_neighbor_match(desc1.unsqueeze(0), desc2.unsqueeze(0))
+
+        matches1 = pred['matches0'][0]
+        mask1 = matches1 != -1
+        pts1 = pts1[mask1, :2]
+        pts2 = pts2[matches1[mask1], :2]
+
+        desc1 = desc1[mask1] # N * 256
+        desc2 = desc2[matches1[mask1]] # N * 256
+
+        K1 = data['view0']['camera'].calibration_matrix()[0] # B * 3 * 3
+        K2 = data['view1']['camera'].calibration_matrix()[0] # B * 3 * 3
+
+        GT_R_Rel = data['T_0to1'].R[0]
+        GT_t_Rel = data['T_0to1'].t[0].unsqueeze(-1)
+
+        ratios = torch.zeros(pts1.shape[0]) # N
+
+        # Patches start
+        bias = torch.tensor([[opt.patch_radius]*2], device=pts1.device)
+        pad_amount = opt.patch_radius
+
+        assert (data['view0']['image'][0] < (1. + 1e-5)).all() and (data['view0']['image'][0] > -(1. + 1e-5)).all(), "Image 0 out of range"
+
+        if data['view0']['image'].shape[1] == 3:  # RGB
+            scale = data['view0']['image'].new_tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1)
+            data['view0']['image'] = (data['view0']['image'] * scale).sum(1, keepdim=True)
+
+            scale = data['view1']['image'].new_tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1)
+            data['view1']['image'] = (data['view1']['image'] * scale).sum(1, keepdim=True)
+
+        img1_padded = torch.nn.functional.pad(data['view0']['image'][0][:, :img1_shape[0], :img1_shape[1]], [pad_amount]*4, mode='constant', value=0.)
+        patch1 = extract_patches(img1_padded, (pts1 - bias + pad_amount).to(device=img1_padded.device, dtype=torch.int32), 2*opt.patch_radius+1)[0]
+
+        assert (data['view1']['image'][0] < (1. + 1e-5)).all() and (data['view1']['image'][0] > -(1. + 1e-5)).all(), "Image 1 out of range"
+        img2_padded = torch.nn.functional.pad(data['view1']['image'][0][:, :img2_shape[0], :img2_shape[1]], [pad_amount]*4, mode='constant', value=0.)
+        patch2 = extract_patches(img2_padded, (pts2 - bias + pad_amount).to(device=img2_padded.device, dtype=torch.int32), 2*opt.patch_radius+1)[0]
+        # Patches end
+
+        with open(out_dir + f'pair-{scene}-{img1_name}-{img2_name}.pkl', 'wb') as f:
+            pickle.dump([
+                pts1.cpu().numpy().astype(np.float32), # N, 2
+                pts2.cpu().numpy().astype(np.float32), # N, 2
+                ratios.cpu().numpy().astype(np.float32), # N
+                img1_shape, # H, W, 3
+                img2_shape, # H, W, 3
+                K1.numpy().astype(np.float32), # 3, 3
+                K2.numpy().astype(np.float32), # 3, 3
+                GT_R_Rel.numpy().astype(np.float32), # 3, 3
+                GT_t_Rel.numpy().astype(np.float32), # 3, 1
+                patch1.cpu().numpy().astype(np.float32), # N, C, 2*patch_radius+1, 2*patch_radius+1
+                patch2.cpu().numpy().astype(np.float32), # N, C, 2*patch_radius+1, 2*patch_radius+1
+                desc1.cpu().numpy().astype(np.float32), # N, D
+                desc2.cpu().numpy().astype(np.float32),  # N, D
+            ], f)
+
+# pts1:  (N, 2)
+# pts2:  (N, 2)
+# ratios:  (N,)
+# img1:  (H, W, 3)
+# img2:  (H, W, 3)
+# K1:  (3, 3)
+# K2:  (3, 3)
+# GT_R_Rel:  (3, 3)
+# GT_t_Rel:  (3, 1)

vismatch/third_party/keypt2subpx/dataset.py

@@ -0,0 +1,145 @@
+import numpy as np
+import torch
+import os
+import cv2
+import math
+import utils
+import pickle
+
+from torch.utils.data import Dataset
+
+class SparsePatchDataset(Dataset):
+	"""Sparse correspondences dataset."""
+
+	def __init__(self, folders, nfeatures, fmat=False, overwrite_side_info=False, without_score=False, without_depth=True, train=False, total_split=-1, current_split=-1):
+		self.nfeatures = nfeatures # ensure fixed number of features, -1 keeps original feature count
+		self.overwrite_side_info = overwrite_side_info # if true, provide no side information to the neural guidance network
+		self.without_score = without_score
+		self.without_depth = without_depth
+		self.train = train
+		
+		# collect precalculated correspondences of all provided datasets
+		self.files = []
+		for folder in folders:
+			self.files += [folder + f for f in os.listdir(folder)]
+		
+		if total_split > 0 and current_split > 0:
+			start_iter = len(self.files) * (current_split-1) // total_split
+			end_iter = len(self.files) * current_split // total_split
+			self.files = self.files[start_iter:end_iter]
+
+		self.fmat = fmat # estimate fundamental matrix instead of essential matrix
+		self.minset = 5 # minimal set size for essential matrices
+		if fmat: self.minset = 7 # minimal set size for fundamental matrices
+			
+	def __len__(self):
+		return len(self.files)
+
+	def __getitem__(self, idx):
+
+		# load precalculated correspondences using pickle
+		data = None
+		with open(self.files[idx], 'rb') as f:
+			data = pickle.load(f)
+		assert data is not None, f"Error: Correspondences {self.files[idx]} could not be loaded."
+
+		# correspondence coordinates and matching ratios (side information)
+		pts1, pts2, ratios = data[0], data[1], data[2] # N x 2, N x 2, N
+		# image sizes
+		im_size1, im_size2 = torch.from_numpy(np.asarray(data[3])), torch.from_numpy(np.asarray(data[4]))
+		# image calibration parameters
+		K1, K2 = torch.from_numpy(data[5]), torch.from_numpy(data[6])
+		# ground truth pose
+		gt_R, gt_t = torch.from_numpy(data[7]), torch.from_numpy(data[8])
+
+		patch1, patch2 = torch.from_numpy(data[9]), torch.from_numpy(data[10])
+		if not self.without_score:
+			scorepatch1, scorepatch2 = torch.from_numpy(data[11]), torch.from_numpy(data[12])
+			descriptor1, descriptor2 = torch.from_numpy(data[13]), torch.from_numpy(data[14])
+		else:
+			descriptor1, descriptor2 = torch.from_numpy(data[11]), torch.from_numpy(data[12])
+		
+		if not self.without_depth:
+			depth1, depth2 = torch.from_numpy(data[-2]), torch.from_numpy(data[-1])
+		
+		if pts1.shape[0] < self.minset:
+			
+			if self.train:
+				print("WARNING! Not enough correspondences left. Only %d correspondences among would be left, so I instead sample another one." % (int(pts1.shape[1])))
+				return self.__getitem__(np.random.randint(0, len(self.files)))
+
+			else:
+				print("WARNING! Not enough correspondences left. Only %d correspondences among would be left, so I instead return Nones." % (int(pts1.shape[1])))
+				nums = 13
+				if not self.without_score:
+					nums += 2
+				if not self.without_depth:
+					nums += 2
+				return [np.zeros(0)]*nums
+		
+		if self.overwrite_side_info:
+			ratios = np.zeros(ratios.shape, dtype=np.float32)
+
+		if self.fmat:
+			# for fundamental matrices, normalize image coordinates using the image size (network should be independent to resolution)
+			utils.normalize_pts(pts1, im_size1)
+			utils.normalize_pts(pts2, im_size2)
+		else:
+			#for essential matrices, normalize image coordinate using the calibration parameters
+			pts1 = cv2.undistortPoints(pts1[None], K1.numpy(), None).squeeze(1)
+			pts2 = cv2.undistortPoints(pts2[None], K2.numpy(), None).squeeze(1)
+
+		# stack image coordinates and side information into one tensor
+		correspondences = np.concatenate((pts1, pts2, ratios[...,None]), axis=-1)
+		# correspondences = np.transpose(correspondences)
+		correspondences = torch.from_numpy(correspondences)
+
+		indices = None
+		if self.nfeatures > 0:
+			# ensure that there are exactly nfeatures entries in the data tensor 
+			if correspondences.size(0) > self.nfeatures:
+				rnd = torch.randperm(correspondences.size(0))
+				correspondences = correspondences[rnd]
+				correspondences = correspondences[0:self.nfeatures]
+				indices = rnd[0:self.nfeatures]
+			else:
+				indices = torch.arange(0, correspondences.size(0))
+
+			if correspondences.size(0) < self.nfeatures:
+				result = correspondences
+				for i in range(0, math.ceil(self.nfeatures / correspondences.size(0) - 1)):
+					rnd = torch.randperm(correspondences.size(0))
+					result = torch.cat((result, correspondences[rnd]), dim=0)
+					indices = torch.cat((indices, rnd))
+				correspondences = result[0:self.nfeatures]
+				indices = indices[0:self.nfeatures]
+
+		# construct the ground truth essential matrix from the ground truth relative pose
+		gt_E = torch.zeros((3,3))
+		gt_E[0, 1] = -float(gt_t[2,0])
+		gt_E[0, 2] = float(gt_t[1,0])
+		gt_E[1, 0] = float(gt_t[2,0])
+		gt_E[1, 2] = -float(gt_t[0,0])
+		gt_E[2, 0] = -float(gt_t[1,0])
+		gt_E[2, 1] = float(gt_t[0,0])
+
+		gt_E = gt_E.mm(gt_R)
+
+		# fundamental matrix from essential matrix
+		gt_F = K2.inverse().transpose(0, 1).mm(gt_E).mm(K1.inverse())
+
+		if indices is not None:
+			patch1, patch2 = patch1[indices], patch2[indices]
+			if not self.without_score:
+				scorepatch1, scorepatch2 = scorepatch1[indices], scorepatch2[indices]
+			descriptor1, descriptor2 = descriptor1[indices], descriptor2[indices]
+			if not self.without_depth:
+				depth1, depth2 = depth1[indices], depth2[indices]
+
+		retlist = [correspondences, gt_F, gt_E, gt_R, gt_t, K1, K2, im_size1, im_size2, patch1, patch2, descriptor1, descriptor2]
+		if not self.without_score:
+			retlist += [scorepatch1, scorepatch2]
+		if not self.without_depth:
+			retlist += [depth1, depth2]
+
+		return retlist

vismatch/third_party/keypt2subpx/hubconf.py

@@ -0,0 +1,38 @@
+dependencies = ['torch']
+from model import Keypt2Subpx as Keypt2Subpx_
+import torch
+
+lentable = {
+    'spnn': 256,
+    'splg': 256,
+    'aliked': 128,
+    'dedode': 256,
+    'xfeat': 64
+}
+
+scoreusetable = {
+    'spnn': True,
+    'splg': True,
+    'aliked': True,
+    'dedode': False,
+    'xfeat': False
+}
+
+def Keypt2Subpx(pretrained=True, detector='splg'):
+    """
+        Keypt2Subpx model
+        pretrained (bool): kwargs, load pretrained weights into the model
+        detector (str): kwargs, detector-matcher combination; supported options: 'spnn', 'splg', 'aliked', 'dedode', 'xfeat'
+    """
+    assert detector in lentable.keys(), f"Unsupported detector-matcher combination; Supported options: {' '.join(lentable.keys())}"
+
+    model = Keypt2Subpx_(lentable[detector], scoreusetable[detector])
+
+    weights = None
+    if pretrained:
+        weights = torch.hub.load_state_dict_from_url(
+            f"https://github.com/KimSinjeong/keypt2subpx/raw/master/pretrained/k2s_{detector}_pretrained.pth",
+            map_location=torch.device('cpu'))
+        model.net.load_state_dict(weights['model'])
+    
+    return model

vismatch/third_party/keypt2subpx/logger.py

@@ -0,0 +1,127 @@
+# A simple torch style logger
+# (C) Wei YANG 2017
+from __future__ import absolute_import
+import matplotlib.pyplot as plt
+import os
+import sys
+import numpy as np
+
+__all__ = ['Logger', 'LoggerMonitor', 'savefig']
+
+def savefig(fname, dpi=None):
+    dpi = 150 if dpi == None else dpi
+    plt.savefig(fname, dpi=dpi)
+    
+def plot_overlap(logger, names=None):
+    names = logger.names if names == None else names
+    numbers = logger.numbers
+    for _, name in enumerate(names):
+        x = np.arange(len(numbers[name]))
+        plt.plot(x, np.asarray(numbers[name]))
+    return [logger.title + '(' + name + ')' for name in names]
+
+class Logger(object):
+    '''Save training process to log file with simple plot function.'''
+    def __init__(self, fpath, title=None, resume=False): 
+        self.file = None
+        self.resume = resume
+        self.title = '' if title == None else title
+        if fpath is not None:
+            if resume: 
+                self.file = open(fpath, 'r') 
+                name = self.file.readline()
+                self.names = name.rstrip().split('\t')
+                self.numbers = {}
+                for _, name in enumerate(self.names):
+                    self.numbers[name] = []
+
+                for numbers in self.file:
+                    numbers = numbers.rstrip().split('\t')
+                    for i in range(0, len(numbers)):
+                        self.numbers[self.names[i]].append(numbers[i])
+                self.file.close()
+                self.file = open(fpath, 'a')  
+            else:
+                self.file = open(fpath, 'w')
+
+    def set_names(self, names):
+        if self.resume: 
+            pass
+        # initialize numbers as empty list
+        self.numbers = {}
+        self.names = names
+        for _, name in enumerate(self.names):
+            self.file.write(name)
+            self.file.write('\t')
+            self.numbers[name] = []
+        self.file.write('\n')
+        self.file.flush()
+
+
+    def append(self, numbers):
+        assert len(self.names) == len(numbers), 'Numbers do not match names'
+        for index, num in enumerate(numbers):
+            self.file.write("{0:.6f}".format(num))
+            self.file.write('\t')
+            self.numbers[self.names[index]].append(num)
+        self.file.write('\n')
+        self.file.flush()
+
+    def plot(self, names=None):   
+        names = self.names if names == None else names
+        numbers = self.numbers
+        for _, name in enumerate(names):
+            x = np.arange(len(numbers[name]))
+            plt.plot(x, np.asarray(numbers[name]))
+        plt.legend([self.title + '(' + name + ')' for name in names])
+        plt.grid(True)
+
+    def close(self):
+        if self.file is not None:
+            self.file.close()
+
+class LoggerMonitor(object):
+    '''Load and visualize multiple logs.'''
+    def __init__ (self, paths):
+        '''paths is a distionary with {name:filepath} pair'''
+        self.loggers = []
+        for title, path in paths.items():
+            logger = Logger(path, title=title, resume=True)
+            self.loggers.append(logger)
+
+    def plot(self, names=None):
+        plt.figure()
+        plt.subplot(121)
+        legend_text = []
+        for logger in self.loggers:
+            legend_text += plot_overlap(logger, names)
+        plt.legend(legend_text, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
+        plt.grid(True)
+                    
+if __name__ == '__main__':
+    # # Example
+    # logger = Logger('test.txt')
+    # logger.set_names(['Train loss', 'Valid loss','Test loss'])
+
+    # length = 100
+    # t = np.arange(length)
+    # train_loss = np.exp(-t / 10.0) + np.random.rand(length) * 0.1
+    # valid_loss = np.exp(-t / 10.0) + np.random.rand(length) * 0.1
+    # test_loss = np.exp(-t / 10.0) + np.random.rand(length) * 0.1
+
+    # for i in range(0, length):
+    #     logger.append([train_loss[i], valid_loss[i], test_loss[i]])
+    # logger.plot()
+
+    # Example: logger monitor
+    paths = {
+    'resadvnet20':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet20/log.txt', 
+    'resadvnet32':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet32/log.txt',
+    'resadvnet44':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet44/log.txt',
+    }
+
+    field = ['Valid Acc.']
+
+    monitor = LoggerMonitor(paths)
+    monitor.plot(names=field)
+    savefig('test.eps')

vismatch/third_party/keypt2subpx/model.py

@@ -0,0 +1,183 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+from utils import extract_patches
+
+def create_meshgrid(
+        x: torch.Tensor,
+        normalized_coordinates: Optional[bool]) -> torch.Tensor:
+    assert len(x.shape) == 4, x.shape
+    _, _, height, width = x.shape
+    _device, _dtype = x.device, x.dtype
+    if normalized_coordinates:
+        xs = torch.linspace(-1.0, 1.0, width, device=_device, dtype=_dtype)
+        ys = torch.linspace(-1.0, 1.0, height, device=_device, dtype=_dtype)
+    else:
+        xs = torch.linspace(0, width - 1, width, device=_device, dtype=_dtype)
+        ys = torch.linspace(0, height - 1, height, device=_device, dtype=_dtype)
+    return torch.meshgrid(ys, xs)  # pos_y, pos_x
+
+class SpatialSoftArgmax2d(nn.Module):
+    r"""Creates a module that computes the Spatial Soft-Argmax 2D
+    of a given input heatmap.
+
+    Returns the index of the maximum 2d coordinates of the give map.
+    The output order is x-coord and y-coord.
+
+    Arguments:
+        normalized_coordinates (Optional[bool]): wether to return the
+          coordinates normalized in the range of [-1, 1]. Otherwise,
+          it will return the coordinates in the range of the input shape.
+          Default is True.
+
+    Shape:
+        - Input: :math:`(B, N, H, W)`
+        - Output: :math:`(B, N, 2)`
+
+    Examples::
+        >>> input = torch.rand(1, 4, 2, 3)
+        >>> m = tgm.losses.SpatialSoftArgmax2d()
+        >>> coords = m(input)  # 1x4x2
+        >>> x_coord, y_coord = torch.chunk(coords, dim=-1, chunks=2)
+    """
+
+    def __init__(self, normalized_coordinates: Optional[bool] = True) -> None:
+        super(SpatialSoftArgmax2d, self).__init__()
+        self.normalized_coordinates: Optional[bool] = normalized_coordinates
+        self.eps: float = 1e-6
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if not torch.is_tensor(input):
+            raise TypeError("Input input type is not a torch.Tensor. Got {}"
+                            .format(type(input)))
+        if not len(input.shape) == 4:
+            raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}"
+                             .format(input.shape))
+        # unpack shapes and create view from input tensor
+        batch_size, channels, height, width = input.shape
+        x: torch.Tensor = input.view(batch_size, channels, -1)
+
+        # compute softmax with max substraction trick
+        exp_x = torch.exp(x - torch.max(x, dim=-1, keepdim=True)[0])
+        exp_x_sum = 1.0 / (exp_x.sum(dim=-1, keepdim=True) + self.eps)
+
+        # create coordinates grid
+        pos_y, pos_x = create_meshgrid(input, self.normalized_coordinates)
+        pos_x = pos_x.reshape(-1)
+        pos_y = pos_y.reshape(-1)
+
+        # compute the expected coordinates
+        expected_y: torch.Tensor = torch.sum(
+            (pos_y * exp_x) * exp_x_sum, dim=-1, keepdim=True)
+        expected_x: torch.Tensor = torch.sum(
+            (pos_x * exp_x) * exp_x_sum, dim=-1, keepdim=True)
+        output: torch.Tensor = torch.cat([expected_x, expected_y], dim=-1)
+        return output.view(batch_size, channels, 2)  # BxNx2
+
+class AttnTuner(nn.Module):
+	def __init__(self, output_dim=256, use_score=True):
+		super(AttnTuner, self).__init__()
+		self.relu = nn.ReLU(inplace=True)
+		self.feat_axis = 1
+		self.normalized_coordinates = False
+		self.use_score = use_score
+
+		c1, c2, c3 = 16, 64, output_dim
+		self.conv1a = nn.Conv2d(2 if use_score else 1, c1, kernel_size=3, stride=1, padding=0)
+		self.conv1b = nn.Conv2d(c1, c1, kernel_size=3, stride=1, padding=1)
+		self.conv2a = nn.Conv2d(c1, c2, kernel_size=3, stride=1, padding=0)
+		self.conv2b = nn.Conv2d(c2, c2, kernel_size=3, stride=1, padding=1)
+		self.conv3 = nn.Conv2d(c2, c3, kernel_size=3, stride=1, padding=0)
+		self.logsoftargmax = SpatialSoftArgmax2d(self.normalized_coordinates)
+
+	def forward(self, patch, scorepatch, desc):
+		B, N, C, H, W = patch.shape
+		B, N, F_ = desc.shape
+		assert H == W, "Patch shape must be square"
+		# P = (H // 2 +1) // 2 +1
+		P = H-6
+
+		patch = patch.view(B*N, C, H, W)
+		if self.use_score:
+			scorepatch = scorepatch.view(B*N, 1, H, W)
+		desc = desc.view(B*N, F_, 1, 1)
+
+		if patch.shape[1] == 3:  # RGB
+			scale = patch.new_tensor([0.299, 0.587, 0.114]).view(*([1]*self.feat_axis), 3, 1, 1)
+			patch = (patch * scale).sum(self.feat_axis, keepdim=True)
+		x = torch.cat([patch, scorepatch], self.feat_axis) if self.use_score else patch
+
+		# Shared Encoder
+		x = self.relu(self.conv1a(x))
+		x = self.relu(self.conv1b(x))
+		x = self.relu(self.conv2a(x))
+		x = self.relu(self.conv2b(x))
+		x = self.conv3(x)
+		x = F.normalize(x, p=2, dim=self.feat_axis)
+		x = (x * desc).sum(dim=self.feat_axis).view(B, N, P, P) # Cosine similarity (in [-1, 1])    
+
+		coord = self.logsoftargmax(x) - (P-1)/2.
+		return coord # B x N x 2
+
+class Keypt2Subpx(nn.Module):
+    def __init__(self, output_dim=256, use_score=True):
+        super(Keypt2Subpx, self).__init__()
+        self.net = AttnTuner(output_dim, use_score)
+        self.use_score = use_score
+        self.patch_radius = 5
+
+    def forward(self, keypt1, keypt2, img1, img2, desc1, desc2, score1=None, score2=None):
+        """
+            keypt1, keypt2: N x 2
+            img1, img2: C x H x W
+            score1, score2: 1 x H x W
+            desc1, desc2: N x D
+            TODO: Batch support
+        """
+        assert (img1 < (1. + 1e-5)).all() and (img1 > -(1. + 1e-5)).all(), "Image 1 out of range"
+        assert (img2 < (1. + 1e-5)).all() and (img2 > -(1. + 1e-5)).all(), "Image 2 out of range"
+        assert desc1 is not None and desc2 is not None
+        C, H, W = img1.shape
+        N, D = desc1.shape
+        assert N == keypt1.shape[0] == keypt2.shape[0], "Number of keypoints mismatch with descriptors"
+
+        bias = torch.tensor([[self.patch_radius]*2], device=keypt1.device)
+
+        # RGB image to grayscale
+        if img1.shape[0] == 3: # RGB
+            scale = img1.new_tensor([0.299, 0.587, 0.114]).view(3, 1, 1)
+            img1 = (img1 * scale).sum(0, keepdim=True)
+            
+            scale = img2.new_tensor([0.299, 0.587, 0.114]).view(3, 1, 1)
+            img2 = (img2 * scale).sum(0, keepdim=True)
+
+        # Image patches
+        img1_padded = torch.nn.functional.pad(img1, [self.patch_radius]*4, mode='constant', value=0.)
+        idx1 = (keypt1 - bias + self.patch_radius).int()
+        patch1 = extract_patches(img1_padded.to(device=idx1.device), idx1, 2*self.patch_radius+1)[0].unsqueeze(0)
+
+        img2_padded = torch.nn.functional.pad(img2, [self.patch_radius]*4, mode='constant', value=0.)
+        idx2 = (keypt2 - bias + self.patch_radius).int()
+        patch2 = extract_patches(img2_padded.to(device=idx2.device), idx2, 2*self.patch_radius+1)[0].unsqueeze(0)
+
+        # Score patches
+        scorepatch1, scorepatch2 = None, None
+        if self.use_score:
+            score1_padded = torch.nn.functional.pad(score1, [self.patch_radius]*4, mode='constant', value=0.)
+            scorepatch1 = extract_patches(score1_padded.to(device=idx1.device), idx1, 2*self.patch_radius+1)[0].unsqueeze(0)
+
+            score2_padded = torch.nn.functional.pad(score2, [self.patch_radius]*4, mode='constant', value=0.)
+            scorepatch2 = extract_patches(score2_padded.to(device=idx2.device), idx2, 2*self.patch_radius+1)[0].unsqueeze(0)
+        
+        meanft = ((desc1 + desc2) / 2.).unsqueeze(0)
+
+        coord1 = 2.5* self.net(patch1, scorepatch1, meanft).unsqueeze(-1) # 1 x N x 2 x 1
+        # coord1 = 2.5* K1[:2,:2].view(1, 1, 2, 2).cuda().inverse() @ coord1 # 1 x N x 2 x 1
+        coord1 = coord1.view(N, 2)
+
+        coord2 = 2.5* self.net(patch2, scorepatch2, meanft).unsqueeze(-1) # 1 x N x 2 x 1
+        # coord2 = 2.5* K2[:2,:2].view(1, 1, 2, 2).cuda().inverse() @ coord2 # 1 x N x 2 x 1
+        coord2 = coord2.view(N, 2)
+
+        return keypt1 + coord1, keypt2 + coord2