vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/accelerated_features/third_party/ALIKE/demo.py

@@ -0,0 +1,167 @@
+import copy
+import os
+import cv2
+import glob
+import logging
+import argparse
+import numpy as np
+from tqdm import tqdm
+from alike import ALike, configs
+
+
+class ImageLoader(object):
+    def __init__(self, filepath: str):
+        self.N = 3000
+        if filepath.startswith('camera'):
+            camera = int(filepath[6:])
+            self.cap = cv2.VideoCapture(camera)
+            if not self.cap.isOpened():
+                raise IOError(f"Can't open camera {camera}!")
+            logging.info(f'Opened camera {camera}')
+            self.mode = 'camera'
+        elif os.path.exists(filepath):
+            if os.path.isfile(filepath):
+                self.cap = cv2.VideoCapture(filepath)
+                if not self.cap.isOpened():
+                    raise IOError(f"Can't open video {filepath}!")
+                rate = self.cap.get(cv2.CAP_PROP_FPS)
+                self.N = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 1
+                duration = self.N / rate
+                logging.info(f'Opened video {filepath}')
+                logging.info(f'Frames: {self.N}, FPS: {rate}, Duration: {duration}s')
+                self.mode = 'video'
+            else:
+                self.images = glob.glob(os.path.join(filepath, '*.png')) + \
+                              glob.glob(os.path.join(filepath, '*.jpg')) + \
+                              glob.glob(os.path.join(filepath, '*.ppm'))
+                self.images.sort()
+                self.N = len(self.images)
+                logging.info(f'Loading {self.N} images')
+                self.mode = 'images'
+        else:
+            raise IOError('Error filepath (camerax/path of images/path of videos): ', filepath)
+
+    def __getitem__(self, item):
+        if self.mode == 'camera' or self.mode == 'video':
+            if item > self.N:
+                return None
+            ret, img = self.cap.read()
+            if not ret:
+                raise "Can't read image from camera"
+            if self.mode == 'video':
+                self.cap.set(cv2.CAP_PROP_POS_FRAMES, item)
+        elif self.mode == 'images':
+            filename = self.images[item]
+            img = cv2.imread(filename)
+            if img is None:
+                raise Exception('Error reading image %s' % filename)        
+        return img
+
+    def __len__(self):
+        return self.N
+
+
+class SimpleTracker(object):
+    def __init__(self):
+        self.pts_prev = None
+        self.desc_prev = None
+
+    def update(self, img, pts, desc):
+        N_matches = 0
+        if self.pts_prev is None:
+            self.pts_prev = pts
+            self.desc_prev = desc
+
+            out = copy.deepcopy(img)
+            for pt1 in pts:
+                p1 = (int(round(pt1[0])), int(round(pt1[1])))
+                cv2.circle(out, p1, 1, (0, 0, 255), -1, lineType=16)
+        else:
+            matches = self.mnn_mather(self.desc_prev, desc)
+            mpts1, mpts2 = self.pts_prev[matches[:, 0]], pts[matches[:, 1]]
+            N_matches = len(matches)
+
+            out = copy.deepcopy(img)
+            for pt1, pt2 in zip(mpts1, mpts2):
+                p1 = (int(round(pt1[0])), int(round(pt1[1])))
+                p2 = (int(round(pt2[0])), int(round(pt2[1])))
+                cv2.line(out, p1, p2, (0, 255, 0), lineType=16)
+                cv2.circle(out, p2, 1, (0, 0, 255), -1, lineType=16)
+
+            self.pts_prev = pts
+            self.desc_prev = desc
+
+        return out, N_matches
+
+    def mnn_mather(self, desc1, desc2):
+        sim = desc1 @ desc2.transpose()
+        sim[sim < 0.9] = 0
+        nn12 = np.argmax(sim, axis=1)
+        nn21 = np.argmax(sim, axis=0)
+        ids1 = np.arange(0, sim.shape[0])
+        mask = (ids1 == nn21[nn12])
+        matches = np.stack([ids1[mask], nn12[mask]])
+        return matches.transpose()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='ALike Demo.')
+    parser.add_argument('input', type=str, default='',
+                        help='Image directory or movie file or "camera0" (for webcam0).')
+    parser.add_argument('--model', choices=['alike-t', 'alike-s', 'alike-n', 'alike-l'], default="alike-t",
+                        help="The model configuration")
+    parser.add_argument('--device', type=str, default='cuda', help="Running device (default: cuda).")
+    parser.add_argument('--top_k', type=int, default=-1,
+                        help='Detect top K keypoints. -1 for threshold based mode, >0 for top K mode. (default: -1)')
+    parser.add_argument('--scores_th', type=float, default=0.2,
+                        help='Detector score threshold (default: 0.2).')
+    parser.add_argument('--n_limit', type=int, default=5000,
+                        help='Maximum number of keypoints to be detected (default: 5000).')
+    parser.add_argument('--no_display', action='store_true',
+                        help='Do not display images to screen. Useful if running remotely (default: False).')
+    parser.add_argument('--no_sub_pixel', action='store_true',
+                        help='Do not detect sub-pixel keypoints (default: False).')
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+
+    image_loader = ImageLoader(args.input)
+    model = ALike(**configs[args.model],
+                  device=args.device,
+                  top_k=args.top_k,
+                  scores_th=args.scores_th,
+                  n_limit=args.n_limit)
+    tracker = SimpleTracker()
+
+    if not args.no_display:
+        logging.info("Press 'q' to stop!")
+        cv2.namedWindow(args.model)
+
+    runtime = []
+    progress_bar = tqdm(image_loader)
+    for img in progress_bar:
+        if img is None:
+            break
+        
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        pred = model(img_rgb, sub_pixel=not args.no_sub_pixel)
+        kpts = pred['keypoints']
+        desc = pred['descriptors']
+        runtime.append(pred['time'])
+
+        out, N_matches = tracker.update(img, kpts, desc)
+
+        ave_fps = (1. / np.stack(runtime)).mean()
+        status = f"Fps:{ave_fps:.1f}, Keypoints/Matches: {len(kpts)}/{N_matches}"
+        progress_bar.set_description(status)
+
+        if not args.no_display:
+            cv2.setWindowTitle(args.model, args.model + ': ' + status)
+            cv2.imshow(args.model, out)
+            if cv2.waitKey(1) == ord('q'):
+                break
+
+    logging.info('Finished!')
+    if not args.no_display:
+        logging.info('Press any key to exit!')
+        cv2.waitKey()

vismatch/third_party/accelerated_features/third_party/ALIKE/hseq/eval.py

@@ -0,0 +1,162 @@
+import cv2
+import os
+from tqdm import tqdm
+import torch
+import numpy as np
+from extract import extract_method
+
+use_cuda = torch.cuda.is_available()
+device = torch.device('cuda' if use_cuda else 'cpu')
+
+methods = ['d2', 'lfnet', 'superpoint', 'r2d2', 'aslfeat', 'disk',
+           'alike-n', 'alike-l', 'alike-n-ms', 'alike-l-ms']
+names = ['D2-Net(MS)', 'LF-Net(MS)', 'SuperPoint', 'R2D2(MS)', 'ASLFeat(MS)', 'DISK',
+         'ALike-N', 'ALike-L', 'ALike-N(MS)', 'ALike-L(MS)']
+
+top_k = None
+n_i = 52
+n_v = 56
+cache_dir = 'hseq/cache'
+dataset_path = 'hseq/hpatches-sequences-release'
+
+
+def generate_read_function(method, extension='ppm'):
+    def read_function(seq_name, im_idx):
+        aux = np.load(os.path.join(dataset_path, seq_name, '%d.%s.%s' % (im_idx, extension, method)))
+        if top_k is None:
+            return aux['keypoints'], aux['descriptors']
+        else:
+            assert ('scores' in aux)
+            ids = np.argsort(aux['scores'])[-top_k:]
+            return aux['keypoints'][ids, :], aux['descriptors'][ids, :]
+
+    return read_function
+
+
+def mnn_matcher(descriptors_a, descriptors_b):
+    device = descriptors_a.device
+    sim = descriptors_a @ descriptors_b.t()
+    nn12 = torch.max(sim, dim=1)[1]
+    nn21 = torch.max(sim, dim=0)[1]
+    ids1 = torch.arange(0, sim.shape[0], device=device)
+    mask = (ids1 == nn21[nn12])
+    matches = torch.stack([ids1[mask], nn12[mask]])
+    return matches.t().data.cpu().numpy()
+
+
+def homo_trans(coord, H):
+    kpt_num = coord.shape[0]
+    homo_coord = np.concatenate((coord, np.ones((kpt_num, 1))), axis=-1)
+    proj_coord = np.matmul(H, homo_coord.T).T
+    proj_coord = proj_coord / proj_coord[:, 2][..., None]
+    proj_coord = proj_coord[:, 0:2]
+    return proj_coord
+
+
+def benchmark_features(read_feats):
+    lim = [1, 5]
+    rng = np.arange(lim[0], lim[1] + 1)
+
+    seq_names = sorted(os.listdir(dataset_path))
+
+    n_feats = []
+    n_matches = []
+    seq_type = []
+    i_err = {thr: 0 for thr in rng}
+    v_err = {thr: 0 for thr in rng}
+
+    i_err_homo = {thr: 0 for thr in rng}
+    v_err_homo = {thr: 0 for thr in rng}
+
+    for seq_idx, seq_name in tqdm(enumerate(seq_names), total=len(seq_names)):
+        keypoints_a, descriptors_a = read_feats(seq_name, 1)
+        n_feats.append(keypoints_a.shape[0])
+
+        # =========== compute homography
+        ref_img = cv2.imread(os.path.join(dataset_path, seq_name, '1.ppm'))
+        ref_img_shape = ref_img.shape
+
+        for im_idx in range(2, 7):
+            keypoints_b, descriptors_b = read_feats(seq_name, im_idx)
+            n_feats.append(keypoints_b.shape[0])
+
+            matches = mnn_matcher(
+                torch.from_numpy(descriptors_a).to(device=device),
+                torch.from_numpy(descriptors_b).to(device=device)
+            )
+
+            homography = np.loadtxt(os.path.join(dataset_path, seq_name, "H_1_" + str(im_idx)))
+
+            pos_a = keypoints_a[matches[:, 0], : 2]
+            pos_a_h = np.concatenate([pos_a, np.ones([matches.shape[0], 1])], axis=1)
+            pos_b_proj_h = np.transpose(np.dot(homography, np.transpose(pos_a_h)))
+            pos_b_proj = pos_b_proj_h[:, : 2] / pos_b_proj_h[:, 2:]
+
+            pos_b = keypoints_b[matches[:, 1], : 2]
+
+            dist = np.sqrt(np.sum((pos_b - pos_b_proj) ** 2, axis=1))
+
+            n_matches.append(matches.shape[0])
+            seq_type.append(seq_name[0])
+
+            if dist.shape[0] == 0:
+                dist = np.array([float("inf")])
+
+            for thr in rng:
+                if seq_name[0] == 'i':
+                    i_err[thr] += np.mean(dist <= thr)
+                else:
+                    v_err[thr] += np.mean(dist <= thr)
+
+            # =========== compute homography
+            gt_homo = homography
+            pred_homo, _ = cv2.findHomography(keypoints_a[matches[:, 0], : 2], keypoints_b[matches[:, 1], : 2],
+                                              cv2.RANSAC)
+            if pred_homo is None:
+                homo_dist = np.array([float("inf")])
+            else:
+                corners = np.array([[0, 0],
+                                    [ref_img_shape[1] - 1, 0],
+                                    [0, ref_img_shape[0] - 1],
+                                    [ref_img_shape[1] - 1, ref_img_shape[0] - 1]])
+                real_warped_corners = homo_trans(corners, gt_homo)
+                warped_corners = homo_trans(corners, pred_homo)
+                homo_dist = np.mean(np.linalg.norm(real_warped_corners - warped_corners, axis=1))
+
+            for thr in rng:
+                if seq_name[0] == 'i':
+                    i_err_homo[thr] += np.mean(homo_dist <= thr)
+                else:
+                    v_err_homo[thr] += np.mean(homo_dist <= thr)
+
+    seq_type = np.array(seq_type)
+    n_feats = np.array(n_feats)
+    n_matches = np.array(n_matches)
+
+    return i_err, v_err, i_err_homo, v_err_homo, [seq_type, n_feats, n_matches]
+
+
+if __name__ == '__main__':
+    errors = {}
+    for method in methods:
+        output_file = os.path.join(cache_dir, method + '.npy')
+        read_function = generate_read_function(method)
+        if os.path.exists(output_file):
+            errors[method] = np.load(output_file, allow_pickle=True)
+        else:
+            extract_method(method)
+            errors[method] = benchmark_features(read_function)
+            np.save(output_file, errors[method])
+
+    for name, method in zip(names, methods):
+        i_err, v_err, i_err_hom, v_err_hom, _ = errors[method]
+
+        print(f"====={name}=====")
+        print(f"MMA@1 MMA@2 MMA@3 MHA@1 MHA@2 MHA@3: ", end='')
+        for thr in range(1, 4):
+            err = (i_err[thr] + v_err[thr]) / ((n_i + n_v) * 5)
+            print(f"{err * 100:.2f}%", end=' ')
+        for thr in range(1, 4):
+            err_hom = (i_err_hom[thr] + v_err_hom[thr]) / ((n_i + n_v) * 5)
+            print(f"{err_hom * 100:.2f}%", end=' ')
+        print('')

vismatch/third_party/accelerated_features/third_party/ALIKE/hseq/extract.py

@@ -0,0 +1,159 @@
+import os
+import sys
+import cv2
+from pathlib import Path
+import numpy as np
+import torch
+import torch.utils.data as data
+from tqdm import tqdm
+from copy import deepcopy
+from torchvision.transforms import ToTensor
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+from alike import ALike, configs
+
+dataset_root = 'hseq/hpatches-sequences-release'
+use_cuda = torch.cuda.is_available()
+device = 'cuda' if use_cuda else 'cpu'
+methods = ['alike-n', 'alike-l', 'alike-n-ms', 'alike-l-ms']
+
+
+class HPatchesDataset(data.Dataset):
+    def __init__(self, root: str = dataset_root, alteration: str = 'all'):
+        """
+        Args:
+            root: dataset root path
+            alteration: # 'all', 'i' for illumination or 'v' for viewpoint
+        """
+        assert (Path(root).exists()), f"Dataset root path {root} dose not exist!"
+        self.root = root
+
+        # get all image file name
+        self.image0_list = []
+        self.image1_list = []
+        self.homographies = []
+        folders = [x for x in Path(self.root).iterdir() if x.is_dir()]
+        self.seqs = []
+        for folder in folders:
+            if alteration == 'i' and folder.stem[0] != 'i':
+                continue
+            if alteration == 'v' and folder.stem[0] != 'v':
+                continue
+
+            self.seqs.append(folder)
+
+        self.len = len(self.seqs)
+        assert (self.len > 0), f'Can not find PatchDataset in path {self.root}'
+
+    def __getitem__(self, item):
+        folder = self.seqs[item]
+
+        imgs = []
+        homos = []
+        for i in range(1, 7):
+            img = cv2.imread(str(folder / f'{i}.ppm'), cv2.IMREAD_COLOR)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # HxWxC
+            imgs.append(img)
+
+            if i != 1:
+                homo = np.loadtxt(str(folder / f'H_1_{i}')).astype('float32')
+                homos.append(homo)
+
+        return imgs, homos, folder.stem
+
+    def __len__(self):
+        return self.len
+
+    def name(self):
+        return self.__class__
+
+
+def extract_multiscale(model, img, scale_f=2 ** 0.5,
+                       min_scale=1., max_scale=1.,
+                       min_size=0., max_size=99999.,
+                       image_size_max=99999,
+                       n_k=0, sort=False):
+    H_, W_, three = img.shape
+    assert three == 3, "input image shape should be [HxWx3]"
+
+    old_bm = torch.backends.cudnn.benchmark
+    torch.backends.cudnn.benchmark = False  # speedup
+
+    # ==================== image size constraint
+    image = deepcopy(img)
+    max_hw = max(H_, W_)
+    if max_hw > image_size_max:
+        ratio = float(image_size_max / max_hw)
+        image = cv2.resize(image, dsize=None, fx=ratio, fy=ratio)
+
+    # ==================== convert image to tensor
+    H, W, three = image.shape
+    image = ToTensor()(image).unsqueeze(0)
+    image = image.to(device)
+
+    s = 1.0  # current scale factor
+    keypoints, descriptors, scores, scores_maps, descriptor_maps = [], [], [], [], []
+    while s + 0.001 >= max(min_scale, min_size / max(H, W)):
+        if s - 0.001 <= min(max_scale, max_size / max(H, W)):
+            nh, nw = image.shape[2:]
+
+            # extract descriptors
+            with torch.no_grad():
+                descriptor_map, scores_map = model.extract_dense_map(image)
+                keypoints_, descriptors_, scores_, _ = model.dkd(scores_map, descriptor_map)
+
+            keypoints.append(keypoints_[0])
+            descriptors.append(descriptors_[0])
+            scores.append(scores_[0])
+
+        s /= scale_f
+
+        # down-scale the image for next iteration
+        nh, nw = round(H * s), round(W * s)
+        image = torch.nn.functional.interpolate(image, (nh, nw), mode='bilinear', align_corners=False)
+
+    # restore value
+    torch.backends.cudnn.benchmark = old_bm
+
+    keypoints = torch.cat(keypoints)
+    descriptors = torch.cat(descriptors)
+    scores = torch.cat(scores)
+    keypoints = (keypoints + 1) / 2 * keypoints.new_tensor([[W_ - 1, H_ - 1]])
+
+    if sort or 0 < n_k < len(keypoints):
+        indices = torch.argsort(scores, descending=True)
+        keypoints = keypoints[indices]
+        descriptors = descriptors[indices]
+        scores = scores[indices]
+
+    if 0 < n_k < len(keypoints):
+        keypoints = keypoints[0:n_k]
+        descriptors = descriptors[0:n_k]
+        scores = scores[0:n_k]
+
+    return {'keypoints': keypoints, 'descriptors': descriptors, 'scores': scores}
+
+
+def extract_method(m):
+    hpatches = HPatchesDataset(root=dataset_root, alteration='all')
+    model = m[:7]
+    min_scale = 0.3 if m[8:] == 'ms' else 1.0
+
+    model = ALike(**configs[model], device=device, top_k=0, scores_th=0.2, n_limit=5000)
+
+    progbar = tqdm(hpatches, desc='Extracting for {}'.format(m))
+    for imgs, homos, seq_name in progbar:
+        for i in range(1, 7):
+            img = imgs[i - 1]
+            pred = extract_multiscale(model, img, min_scale=min_scale, max_scale=1, sort=False, n_k=5000)
+            kpts, descs, scores = pred['keypoints'], pred['descriptors'], pred['scores']
+
+            with open(os.path.join(dataset_root, seq_name, f'{i}.ppm.{m}'), 'wb') as f:
+                np.savez(f, keypoints=kpts.cpu().numpy(),
+                         scores=scores.cpu().numpy(),
+                         descriptors=descs.cpu().numpy())
+
+
+if __name__ == '__main__':
+    for method in methods:
+        extract_method(method)

vismatch/third_party/accelerated_features/third_party/ALIKE/soft_detect.py

@@ -0,0 +1,194 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+# coordinates system
+#  ------------------------------>  [ x: range=-1.0~1.0; w: range=0~W ]
+#  | -----------------------------
+#  | |                           |
+#  | |                           |
+#  | |                           |
+#  | |         image             |
+#  | |                           |
+#  | |                           |
+#  | |                           |
+#  | |---------------------------|
+#  v
+# [ y: range=-1.0~1.0; h: range=0~H ]
+
+def simple_nms(scores, nms_radius: int):
+    """ Fast Non-maximum suppression to remove nearby points """
+    assert (nms_radius >= 0)
+
+    def max_pool(x):
+        return torch.nn.functional.max_pool2d(
+            x, kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius)
+
+    zeros = torch.zeros_like(scores)
+    max_mask = scores == max_pool(scores)
+
+    for _ in range(2):
+        supp_mask = max_pool(max_mask.float()) > 0
+        supp_scores = torch.where(supp_mask, zeros, scores)
+        new_max_mask = supp_scores == max_pool(supp_scores)
+        max_mask = max_mask | (new_max_mask & (~supp_mask))
+    return torch.where(max_mask, scores, zeros)
+
+
+def sample_descriptor(descriptor_map, kpts, bilinear_interp=False):
+    """
+    :param descriptor_map: BxCxHxW
+    :param kpts: list, len=B, each is Nx2 (keypoints) [h,w]
+    :param bilinear_interp: bool, whether to use bilinear interpolation
+    :return: descriptors: list, len=B, each is NxD
+    """
+    batch_size, channel, height, width = descriptor_map.shape
+
+    descriptors = []
+    for index in range(batch_size):
+        kptsi = kpts[index]  # Nx2,(x,y)
+
+        if bilinear_interp:
+            descriptors_ = torch.nn.functional.grid_sample(descriptor_map[index].unsqueeze(0), kptsi.view(1, 1, -1, 2),
+                                                           mode='bilinear', align_corners=True)[0, :, 0, :]  # CxN
+        else:
+            kptsi = (kptsi + 1) / 2 * kptsi.new_tensor([[width - 1, height - 1]])
+            kptsi = kptsi.long()
+            descriptors_ = descriptor_map[index, :, kptsi[:, 1], kptsi[:, 0]]  # CxN
+
+        descriptors_ = torch.nn.functional.normalize(descriptors_, p=2, dim=0)
+        descriptors.append(descriptors_.t())
+
+    return descriptors
+
+
+class DKD(nn.Module):
+    def __init__(self, radius=2, top_k=0, scores_th=0.2, n_limit=20000):
+        """
+        Args:
+            radius: soft detection radius, kernel size is (2 * radius + 1)
+            top_k: top_k > 0: return top k keypoints
+            scores_th: top_k <= 0 threshold mode:  scores_th > 0: return keypoints with scores>scores_th
+                                                   else: return keypoints with scores > scores.mean()
+            n_limit: max number of keypoint in threshold mode
+        """
+        super().__init__()
+        self.radius = radius
+        self.top_k = top_k
+        self.scores_th = scores_th
+        self.n_limit = n_limit
+        self.kernel_size = 2 * self.radius + 1
+        self.temperature = 0.1  # tuned temperature
+        self.unfold = nn.Unfold(kernel_size=self.kernel_size, padding=self.radius)
+
+        # local xy grid
+        x = torch.linspace(-self.radius, self.radius, self.kernel_size)
+        # (kernel_size*kernel_size) x 2 : (w,h)
+        self.hw_grid = torch.stack(torch.meshgrid([x, x])).view(2, -1).t()[:, [1, 0]]
+
+    def detect_keypoints(self, scores_map, sub_pixel=True):
+        b, c, h, w = scores_map.shape
+        scores_nograd = scores_map.detach()
+        # nms_scores = simple_nms(scores_nograd, self.radius)
+        nms_scores = simple_nms(scores_nograd, 2)
+
+        # remove border
+        nms_scores[:, :, :self.radius + 1, :] = 0
+        nms_scores[:, :, :, :self.radius + 1] = 0
+        nms_scores[:, :, h - self.radius:, :] = 0
+        nms_scores[:, :, :, w - self.radius:] = 0
+
+        # detect keypoints without grad
+        if self.top_k > 0:
+            topk = torch.topk(nms_scores.view(b, -1), self.top_k)
+            indices_keypoints = topk.indices  # B x top_k
+        else:
+            if self.scores_th > 0:
+                masks = nms_scores > self.scores_th
+                if masks.sum() == 0:
+                    th = scores_nograd.reshape(b, -1).mean(dim=1)  # th = self.scores_th
+                    masks = nms_scores > th.reshape(b, 1, 1, 1)
+            else:
+                th = scores_nograd.reshape(b, -1).mean(dim=1)  # th = self.scores_th
+                masks = nms_scores > th.reshape(b, 1, 1, 1)
+            masks = masks.reshape(b, -1)
+
+            indices_keypoints = []  # list, B x (any size)
+            scores_view = scores_nograd.reshape(b, -1)
+            for mask, scores in zip(masks, scores_view):
+                indices = mask.nonzero(as_tuple=False)[:, 0]
+                if len(indices) > self.n_limit:
+                    kpts_sc = scores[indices]
+                    sort_idx = kpts_sc.sort(descending=True)[1]
+                    sel_idx = sort_idx[:self.n_limit]
+                    indices = indices[sel_idx]
+                indices_keypoints.append(indices)
+
+        keypoints = []
+        scoredispersitys = []
+        kptscores = []
+        if sub_pixel:
+            # detect soft keypoints with grad backpropagation
+            patches = self.unfold(scores_map)  # B x (kernel**2) x (H*W)
+            self.hw_grid = self.hw_grid.to(patches)  # to device
+            for b_idx in range(b):
+                patch = patches[b_idx].t()  # (H*W) x (kernel**2)
+                indices_kpt = indices_keypoints[b_idx]  # one dimension vector, say its size is M
+                patch_scores = patch[indices_kpt]  # M x (kernel**2)
+
+                # max is detached to prevent undesired backprop loops in the graph
+                max_v = patch_scores.max(dim=1).values.detach()[:, None]
+                x_exp = ((patch_scores - max_v) / self.temperature).exp()  # M * (kernel**2), in [0, 1]
+
+                # \frac{ \sum{(i,j) \times \exp(x/T)} }{ \sum{\exp(x/T)} }
+                xy_residual = x_exp @ self.hw_grid / x_exp.sum(dim=1)[:, None]  # Soft-argmax, Mx2
+
+                hw_grid_dist2 = torch.norm((self.hw_grid[None, :, :] - xy_residual[:, None, :]) / self.radius,
+                                           dim=-1) ** 2
+                scoredispersity = (x_exp * hw_grid_dist2).sum(dim=1) / x_exp.sum(dim=1)
+
+                # compute result keypoints
+                keypoints_xy_nms = torch.stack([indices_kpt % w, indices_kpt // w], dim=1)  # Mx2
+                keypoints_xy = keypoints_xy_nms + xy_residual
+                keypoints_xy = keypoints_xy / keypoints_xy.new_tensor(
+                    [w - 1, h - 1]) * 2 - 1  # (w,h) -> (-1~1,-1~1)
+
+                kptscore = torch.nn.functional.grid_sample(scores_map[b_idx].unsqueeze(0),
+                                                           keypoints_xy.view(1, 1, -1, 2),
+                                                           mode='bilinear', align_corners=True)[0, 0, 0, :]  # CxN
+
+                keypoints.append(keypoints_xy)
+                scoredispersitys.append(scoredispersity)
+                kptscores.append(kptscore)
+        else:
+            for b_idx in range(b):
+                indices_kpt = indices_keypoints[b_idx]  # one dimension vector, say its size is M
+                keypoints_xy_nms = torch.stack([indices_kpt % w, indices_kpt // w], dim=1)  # Mx2
+                keypoints_xy = keypoints_xy_nms / keypoints_xy_nms.new_tensor(
+                    [w - 1, h - 1]) * 2 - 1  # (w,h) -> (-1~1,-1~1)
+                kptscore = torch.nn.functional.grid_sample(scores_map[b_idx].unsqueeze(0),
+                                                           keypoints_xy.view(1, 1, -1, 2),
+                                                           mode='bilinear', align_corners=True)[0, 0, 0, :]  # CxN
+                keypoints.append(keypoints_xy)
+                scoredispersitys.append(None)
+                kptscores.append(kptscore)
+
+        return keypoints, scoredispersitys, kptscores
+
+    def forward(self, scores_map, descriptor_map, sub_pixel=False):
+        """
+        :param scores_map:  Bx1xHxW
+        :param descriptor_map: BxCxHxW
+        :param sub_pixel: whether to use sub-pixel keypoint detection
+        :return: kpts: list[Nx2,...]; kptscores: list[N,....] normalised position: -1.0 ~ 1.0
+        """
+        keypoints, scoredispersitys, kptscores = self.detect_keypoints(scores_map,
+                                                                       sub_pixel)
+
+        descriptors = sample_descriptor(descriptor_map, keypoints, sub_pixel)
+
+        # keypoints: B M 2
+        # descriptors: B M D
+        # scoredispersitys:
+        return keypoints, descriptors, kptscores, scoredispersitys

vismatch/third_party/accelerated_features/third_party/__init__.py

@@ -0,0 +1,4 @@
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+"""