add: thirdparty

third_party/ALIKE/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2022, Zhao Xiaoming
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

third_party/ALIKE/README.md

@@ -0,0 +1,131 @@
+# News
+
+- The [ALIKED](https://github.com/Shiaoming/ALIKED) is released.
+- The [ALIKE training code](https://github.com/Shiaoming/ALIKE/raw/main/assets/ALIKE_code.zip) is released.
+
+# ALIKE: Accurate and Lightweight Keypoint Detection and Descriptor Extraction
+
+ALIKE applies a differentiable keypoint detection module to detect accurate sub-pixel keypoints. The network can run at 95 frames per second for 640 x 480 images on NVIDIA Titan X (Pascal) GPU and achieve equivalent performance with the state-of-the-arts. ALIKE benefits real-time applications in resource-limited platforms/devices. Technical details are described in [this paper](https://arxiv.org/pdf/2112.02906.pdf).
+
+> ```
+> Xiaoming Zhao, Xingming Wu, Jinyu Miao, Weihai Chen, Peter C. Y. Chen, Zhengguo Li, "ALIKE: Accurate and Lightweight Keypoint
+> Detection and Descriptor Extraction," IEEE Transactions on Multimedia, 2022.
+> ```
+
+![](./assets/alike.png)
+
+
+If you use ALIKE in an academic work, please cite:
+
+```
+@article{Zhao2023ALIKED,
+    title = {ALIKED: A Lighter Keypoint and Descriptor Extraction Network via Deformable Transformation},
+    url = {https://arxiv.org/pdf/2304.03608.pdf},
+    doi = {10.1109/TIM.2023.3271000},
+    journal = {IEEE Transactions on Instrumentation & Measurement},
+    author = {Zhao, Xiaoming and Wu, Xingming and Chen, Weihai and Chen, Peter C. Y. and Xu, Qingsong and Li, Zhengguo},
+    year = {2023},
+    volume = {72},
+    pages = {1-16},
+}
+
+@article{Zhao2022ALIKE,
+    title = {ALIKE: Accurate and Lightweight Keypoint Detection and Descriptor Extraction},
+    url = {http://arxiv.org/abs/2112.02906},
+    doi = {10.1109/TMM.2022.3155927},
+    journal = {IEEE Transactions on Multimedia},
+    author = {Zhao, Xiaoming and Wu, Xingming and Miao, Jinyu and Chen, Weihai and Chen, Peter C. Y. and Li, Zhengguo},
+    month = march,
+    year = {2022},
+}
+```
+
+
+
+## 1. Prerequisites
+
+The required packages are listed in the `requirements.txt` :
+
+```shell
+pip install -r requirements.txt
+```
+
+
+
+## 2. Models
+
+The off-the-shelf weights of four variant ALIKE models are provided in `models/` .
+
+
+
+## 3. Run demo
+
+```shell
+$ python demo.py -h
+usage: demo.py [-h] [--model {alike-t,alike-s,alike-n,alike-l}]
+               [--device DEVICE] [--top_k TOP_K] [--scores_th SCORES_TH]
+               [--n_limit N_LIMIT] [--no_display] [--no_sub_pixel]
+               input
+
+ALike Demo.
+
+positional arguments:
+  input                 Image directory or movie file or "camera0" (for
+                        webcam0).
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --model {alike-t,alike-s,alike-n,alike-l}
+                        The model configuration
+  --device DEVICE       Running device (default: cuda).
+  --top_k TOP_K         Detect top K keypoints. -1 for threshold based mode,
+                        >0 for top K mode. (default: -1)
+  --scores_th SCORES_TH
+                        Detector score threshold (default: 0.2).
+  --n_limit N_LIMIT     Maximum number of keypoints to be detected (default:
+                        5000).
+  --no_display          Do not display images to screen. Useful if running
+                        remotely (default: False).
+  --no_sub_pixel        Do not detect sub-pixel keypoints (default: False).
+```
+
+
+
+## 4. Examples
+
+### KITTI example
+```shell
+python demo.py assets/kitti 
+```
+![](./assets/kitti.gif)
+
+### TUM example
+```shell
+python demo.py assets/tum 
+```
+![](./assets/tum.gif)
+
+## 5. Efficiency and performance
+
+| Models | Parameters | GFLOPs(640x480) | MHA@3 on Hpatches | mAA(10°) on [IMW2020-test](https://www.cs.ubc.ca/research/image-matching-challenge/2021/leaderboard) (Stereo) |
+|:---:|:---:|:---:|:-----------------:|:-------------------------------------------------------------------------------------------------------------:|
+| D2-Net(MS) | 7653KB | 889.40 |      38.33%       |                                                    12.27%                                                     |
+| LF-Net(MS) | 2642KB | 24.37 |      57.78%       |                                                    23.44%                                                     |
+| SuperPoint | 1301KB | 26.11 |      70.19%       |                                                    28.97%                                                     |
+| R2D2(MS) | 484KB | 464.55 |      71.48%       |                                                    39.02%                                                     |
+| ASLFeat(MS) | 823KB | 77.58 |      73.52%       |                                                    33.65%                                                     |
+| DISK | 1092KB | 98.97 |      70.56%       |                                                    51.22%                                                     |
+| ALike-N | 318KB | 7.909 |      75.74%       |                                                    47.18%                                                     |
+| ALike-L | 653KB | 19.685 |      76.85%       |                                                    49.58%                                                     |
+
+### Evaluation on Hpatches
+
+- Download [hpatches-sequences-release](https://hpatches.github.io/) and put it into `hseq/hpatches-sequences-release`.
+- Remove the unreliable sequences as D2-Net.
+- Run the following command to evaluate the performance:
+  ```shell  
+  python hseq/eval.py
+  ```
+
+
+For more details, please refer to the [paper](https://arxiv.org/abs/2112.02906).

third_party/ALIKE/alike.py

@@ -0,0 +1,198 @@
+import logging
+import os
+import cv2
+import torch
+from copy import deepcopy
+import torch.nn.functional as F
+from torchvision.transforms import ToTensor
+import math
+
+from alnet import ALNet
+from soft_detect import DKD
+import time
+
+configs = {
+    "alike-t": {
+        "c1": 8,
+        "c2": 16,
+        "c3": 32,
+        "c4": 64,
+        "dim": 64,
+        "single_head": True,
+        "radius": 2,
+        "model_path": os.path.join(os.path.split(__file__)[0], "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(os.path.split(__file__)[0], "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(os.path.split(__file__)[0], "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(os.path.split(__file__)[0], "models", "alike-l.pth"),
+    },
+}
+
+
+class ALike(ALNet):
+    def __init__(
+        self,
+        # ================================== feature encoder
+        c1: int = 32,
+        c2: int = 64,
+        c3: int = 128,
+        c4: int = 128,
+        dim: int = 128,
+        single_head: bool = False,
+        # ================================== detect parameters
+        radius: int = 2,
+        top_k: int = 500,
+        scores_th: float = 0.5,
+        n_limit: int = 5000,
+        device: str = "cpu",
+        model_path: str = "",
+    ):
+        super().__init__(c1, c2, c3, c4, dim, single_head)
+        self.radius = radius
+        self.top_k = top_k
+        self.n_limit = n_limit
+        self.scores_th = scores_th
+        self.dkd = DKD(
+            radius=self.radius,
+            top_k=self.top_k,
+            scores_th=self.scores_th,
+            n_limit=self.n_limit,
+        )
+        self.device = device
+
+        if model_path != "":
+            state_dict = torch.load(model_path, self.device)
+            self.load_state_dict(state_dict)
+            self.to(self.device)
+            self.eval()
+            logging.info(f"Loaded model parameters from {model_path}")
+            logging.info(
+                f"Number of model parameters: {sum(p.numel() for p in self.parameters() if p.requires_grad) / 1e3}KB"
+            )
+
+    def extract_dense_map(self, image, ret_dict=False):
+        # ====================================================
+        # check image size, should be integer multiples of 2^5
+        # if it is not a integer multiples of 2^5, padding zeros
+        device = image.device
+        b, c, h, w = image.shape
+        h_ = math.ceil(h / 32) * 32 if h % 32 != 0 else h
+        w_ = math.ceil(w / 32) * 32 if w % 32 != 0 else w
+        if h_ != h:
+            h_padding = torch.zeros(b, c, h_ - h, w, device=device)
+            image = torch.cat([image, h_padding], dim=2)
+        if w_ != w:
+            w_padding = torch.zeros(b, c, h_, w_ - w, device=device)
+            image = torch.cat([image, w_padding], dim=3)
+        # ====================================================
+
+        scores_map, descriptor_map = super().forward(image)
+
+        # ====================================================
+        if h_ != h or w_ != w:
+            descriptor_map = descriptor_map[:, :, :h, :w]
+            scores_map = scores_map[:, :, :h, :w]  # Bx1xHxW
+        # ====================================================
+
+        # BxCxHxW
+        descriptor_map = torch.nn.functional.normalize(descriptor_map, p=2, dim=1)
+
+        if ret_dict:
+            return {
+                "descriptor_map": descriptor_map,
+                "scores_map": scores_map,
+            }
+        else:
+            return descriptor_map, scores_map
+
+    def forward(self, img, image_size_max=99999, sort=False, sub_pixel=False):
+        """
+        :param img: np.array HxWx3, RGB
+        :param image_size_max: maximum image size, otherwise, the image will be resized
+        :param sort: sort keypoints by scores
+        :param sub_pixel: whether to use sub-pixel accuracy
+        :return: a dictionary with 'keypoints', 'descriptors', 'scores', and 'time'
+        """
+        H, W, three = img.shape
+        assert three == 3, "input image shape should be [HxWx3]"
+
+        # ==================== 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
+        image = (
+            torch.from_numpy(image)
+            .to(self.device)
+            .to(torch.float32)
+            .permute(2, 0, 1)[None]
+            / 255.0
+        )
+
+        # ==================== extract keypoints
+        start = time.time()
+
+        with torch.no_grad():
+            descriptor_map, scores_map = self.extract_dense_map(image)
+            keypoints, descriptors, scores, _ = self.dkd(
+                scores_map, descriptor_map, sub_pixel=sub_pixel
+            )
+            keypoints, descriptors, scores = keypoints[0], descriptors[0], scores[0]
+            keypoints = (keypoints + 1) / 2 * keypoints.new_tensor([[W - 1, H - 1]])
+
+        if sort:
+            indices = torch.argsort(scores, descending=True)
+            keypoints = keypoints[indices]
+            descriptors = descriptors[indices]
+            scores = scores[indices]
+
+        end = time.time()
+
+        return {
+            "keypoints": keypoints.cpu().numpy(),
+            "descriptors": descriptors.cpu().numpy(),
+            "scores": scores.cpu().numpy(),
+            "scores_map": scores_map.cpu().numpy(),
+            "time": end - start,
+        }
+
+
+if __name__ == "__main__":
+    import numpy as np
+    from thop import profile
+
+    net = ALike(c1=32, c2=64, c3=128, c4=128, dim=128, single_head=False)
+
+    image = np.random.random((640, 480, 3)).astype(np.float32)
+    flops, params = profile(net, inputs=(image, 9999, False), verbose=False)
+    print("{:<30}  {:<8} GFLops".format("Computational complexity: ", flops / 1e9))
+    print("{:<30}  {:<8} KB".format("Number of parameters: ", params / 1e3))

third_party/ALIKE/alnet.py

@@ -0,0 +1,194 @@
+import torch
+from torch import nn
+from torchvision.models import resnet
+from typing import Optional, Callable
+
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        gate: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ):
+        super().__init__()
+        if gate is None:
+            self.gate = nn.ReLU(inplace=True)
+        else:
+            self.gate = gate
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self.conv1 = resnet.conv3x3(in_channels, out_channels)
+        self.bn1 = norm_layer(out_channels)
+        self.conv2 = resnet.conv3x3(out_channels, out_channels)
+        self.bn2 = norm_layer(out_channels)
+
+    def forward(self, x):
+        x = self.gate(self.bn1(self.conv1(x)))  # B x in_channels x H x W
+        x = self.gate(self.bn2(self.conv2(x)))  # B x out_channels x H x W
+        return x
+
+
+# copied from torchvision\models\resnet.py#27->BasicBlock
+class ResBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        gate: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super(ResBlock, self).__init__()
+        if gate is None:
+            self.gate = nn.ReLU(inplace=True)
+        else:
+            self.gate = gate
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError("ResBlock only supports groups=1 and base_width=64")
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in ResBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = resnet.conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.conv2 = resnet.conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.gate(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.gate(out)
+
+        return out
+
+
+class ALNet(nn.Module):
+    def __init__(
+        self,
+        c1: int = 32,
+        c2: int = 64,
+        c3: int = 128,
+        c4: int = 128,
+        dim: int = 128,
+        single_head: bool = True,
+    ):
+        super().__init__()
+
+        self.gate = nn.ReLU(inplace=True)
+
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.pool4 = nn.MaxPool2d(kernel_size=4, stride=4)
+
+        self.block1 = ConvBlock(3, c1, self.gate, nn.BatchNorm2d)
+
+        self.block2 = ResBlock(
+            inplanes=c1,
+            planes=c2,
+            stride=1,
+            downsample=nn.Conv2d(c1, c2, 1),
+            gate=self.gate,
+            norm_layer=nn.BatchNorm2d,
+        )
+        self.block3 = ResBlock(
+            inplanes=c2,
+            planes=c3,
+            stride=1,
+            downsample=nn.Conv2d(c2, c3, 1),
+            gate=self.gate,
+            norm_layer=nn.BatchNorm2d,
+        )
+        self.block4 = ResBlock(
+            inplanes=c3,
+            planes=c4,
+            stride=1,
+            downsample=nn.Conv2d(c3, c4, 1),
+            gate=self.gate,
+            norm_layer=nn.BatchNorm2d,
+        )
+
+        # ================================== feature aggregation
+        self.conv1 = resnet.conv1x1(c1, dim // 4)
+        self.conv2 = resnet.conv1x1(c2, dim // 4)
+        self.conv3 = resnet.conv1x1(c3, dim // 4)
+        self.conv4 = resnet.conv1x1(dim, dim // 4)
+        self.upsample2 = nn.Upsample(
+            scale_factor=2, mode="bilinear", align_corners=True
+        )
+        self.upsample4 = nn.Upsample(
+            scale_factor=4, mode="bilinear", align_corners=True
+        )
+        self.upsample8 = nn.Upsample(
+            scale_factor=8, mode="bilinear", align_corners=True
+        )
+        self.upsample32 = nn.Upsample(
+            scale_factor=32, mode="bilinear", align_corners=True
+        )
+
+        # ================================== detector and descriptor head
+        self.single_head = single_head
+        if not self.single_head:
+            self.convhead1 = resnet.conv1x1(dim, dim)
+        self.convhead2 = resnet.conv1x1(dim, dim + 1)
+
+    def forward(self, image):
+        # ================================== feature encoder
+        x1 = self.block1(image)  # B x c1 x H x W
+        x2 = self.pool2(x1)
+        x2 = self.block2(x2)  # B x c2 x H/2 x W/2
+        x3 = self.pool4(x2)
+        x3 = self.block3(x3)  # B x c3 x H/8 x W/8
+        x4 = self.pool4(x3)
+        x4 = self.block4(x4)  # B x dim x H/32 x W/32
+
+        # ================================== feature aggregation
+        x1 = self.gate(self.conv1(x1))  # B x dim//4 x H x W
+        x2 = self.gate(self.conv2(x2))  # B x dim//4 x H//2 x W//2
+        x3 = self.gate(self.conv3(x3))  # B x dim//4 x H//8 x W//8
+        x4 = self.gate(self.conv4(x4))  # B x dim//4 x H//32 x W//32
+        x2_up = self.upsample2(x2)  # B x dim//4 x H x W
+        x3_up = self.upsample8(x3)  # B x dim//4 x H x W
+        x4_up = self.upsample32(x4)  # B x dim//4 x H x W
+        x1234 = torch.cat([x1, x2_up, x3_up, x4_up], dim=1)
+
+        # ================================== detector and descriptor head
+        if not self.single_head:
+            x1234 = self.gate(self.convhead1(x1234))
+        x = self.convhead2(x1234)  # B x dim+1 x H x W
+
+        descriptor_map = x[:, :-1, :, :]
+        scores_map = torch.sigmoid(x[:, -1, :, :]).unsqueeze(1)
+
+        return scores_map, descriptor_map
+
+
+if __name__ == "__main__":
+    from thop import profile
+
+    net = ALNet(c1=16, c2=32, c3=64, c4=128, dim=128, single_head=True)
+
+    image = torch.randn(1, 3, 640, 480)
+    flops, params = profile(net, inputs=(image,), verbose=False)
+    print("{:<30}  {:<8} GFLops".format("Computational complexity: ", flops / 1e9))
+    print("{:<30}  {:<8} KB".format("Number of parameters: ", params / 1e3))

third_party/ALIKE/demo.py

@@ -0,0 +1,201 @@
+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.0 / 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()

third_party/ALIKE/hseq/cache/alike-l-ms.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1350ab826afdd9b7542a556e2fda9ad9f94388a875c8edb7874e4bcdfebc63ca
+size 13124

third_party/ALIKE/hseq/cache/alike-l.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:999daff1155f3d4736bb7374fb2058f520b0cb4c75b5d7d87fc1e7025a7d2a7d
+size 13124

third_party/ALIKE/hseq/cache/alike-n-ms.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e5967048eddb61e423bf2ea05a2a626e18d8a716b6a0ad42471059aec0b934c
+size 13124

third_party/ALIKE/hseq/cache/alike-n.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e2eba5ff96b25d0a100b6c7273549de91586e6069dcb5320a20edbb24ea462e
+size 13124

third_party/ALIKE/hseq/cache/aslfeat.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce06fd1b6265e09ed3b26768b68f624e2d556358ab98addd8ebdb7a5a076abe8
+size 15352

third_party/ALIKE/hseq/cache/d2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:976d81c6b51a98f89eac60c6d25990130c1df571ef6536280f4b00577eab56f0
+size 15352

third_party/ALIKE/hseq/cache/disk.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df2d9e0dfd0baa19f2af12f4604368ca65a1643159e7e3438e25efc41ab15357
+size 15352

third_party/ALIKE/hseq/cache/lfnet.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:417327dee726cffccc6dfbc9b0e6b3c06b277ea8878ccf87b87475d1cd6e65ca
+size 15352

third_party/ALIKE/hseq/cache/r2d2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1375a21adcc932db2c9e210e52f633c1903cca6d37066391eb9d645ff87d0120
+size 15352

third_party/ALIKE/hseq/cache/superpoint.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e4d4a4ca79518af47467e9ddd69fe159c9305a580dadc4fdab6ffde6f8b48c2
+size 15352

third_party/ALIKE/hseq/eval.py

@@ -0,0 +1,197 @@
+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("")

third_party/ALIKE/hseq/extract.py

@@ -0,0 +1,175 @@
+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.0,
+    max_scale=1.0,
+    min_size=0.0,
+    max_size=99999.0,
+    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)

third_party/ALIKE/matlab/createfigure.m

@@ -0,0 +1,75 @@
+function createfigure(X1, YMatrix1, Y1, l1, l2, l3)
+%CREATEFIGURE(X1, YMatrix1, Y1)
+%  X1:  vector of x data
+%  YMATRIX1:  matrix of y data
+%  Y1:  vector of y data
+
+%  Auto-generated by MATLAB on 29-Oct-2021 15:42:14
+
+% Create figure
+figure1 = figure;
+
+% Create axes
+axes1 = axes('Parent',figure1);
+hold(axes1,'on');
+
+% Create multiple lines using matrix input to plot
+plot1 = plot(X1,YMatrix1,'Parent',axes1,'LineWidth',1);
+set(plot1(1),'LineStyle','-.','Color',[1 0 0]);
+set(plot1(2),'Color',[0 1 0]);
+set(plot1(3),'LineStyle','--',...
+    'Color',[0.87058824300766 0.490196079015732 0]);
+
+% Uncomment the following line to preserve the X-limits of the axes
+% xlim(axes1,[-1.1 1.1]);
+% Uncomment the following line to preserve the Y-limits of the axes
+ylim(axes1,[0 2.2]);
+box(axes1,'on');
+hold(axes1,'off');
+% Set the remaining axes properties
+set(axes1,'XColor',[0 0 0],'YColor',[0 0 0],'YTick',[0 0.5 1 1.5 2 2.5]);
+% Create axes
+axes2 = axes('Parent',figure1);
+hold(axes2,'on');
+colororder([0.494 0.184 0.556;0.466 0.674 0.188;0.301 0.745 0.933;0.635 0.078 0.184;0 0.447 0.741;0.85 0.325 0.098;0.929 0.694 0.125]);
+
+% Create plot
+plot(X1,Y1,'Parent',axes2,'LineWidth',1,'LineStyle',':','Color',[0 0 1]);
+
+% Uncomment the following line to preserve the X-limits of the axes
+% xlim(axes2,[-1.1 1.1]);
+% Uncomment the following line to preserve the Y-limits of the axes
+ylim(axes2,[0 1.6]);
+hold(axes2,'off');
+% Set the remaining axes properties
+set(axes2,'Color','none','HitTest','off','XColor',[0 0 0],'YAxisLocation',...
+    'right','YColor',[0 0 0],'YTick',[0 0.5 1 1.5]);
+% Create textbox
+annotation(figure1,'textbox',...
+    [0.255427607968038,0.605539475745798,0.304947448327989,0.235148519909872],...
+    'Color',[0.8 0 0],...
+    'String',{sprintf('peak loss=%.4f',l1)},...
+    'EdgeColor','none');
+
+% Create textbox
+annotation(figure1,'textbox',...
+    [0.631790371410027,0.083530640355914,0.178879315581032,0.235148519909871],...
+    'Color',[0 0 1],...
+    'String',{'keypoint'},...
+    'EdgeColor','none');
+
+% Create textbox
+annotation(figure1,'textbox',...
+    [0.59663112557549,0.640686239621974,0.318247136419826,0.22093023731067],...
+    'Color',[0 0.498039215803146 0],...
+    'String',{sprintf('peak loss=%.4f',l2)},...
+    'EdgeColor','none');
+
+% Create textbox
+annotation(figure1,'textbox',...
+    [0.595423071596731,0.415858983920567,0.318247136419826,0.235148519909871],...
+    'Color',[0.87058824300766 0.490196079015732 0],...
+    'String',{sprintf('peak loss=%.4f',l3)},...
+    'FitBoxToText','off',...
+    'EdgeColor','none');
+

third_party/ALIKE/matlab/peakloss_rect.m

@@ -0,0 +1,19 @@
+clear;
+close all;
+
+x = -1:0.01:1;
+
+p0 = 0.5;
+p1 = -0.5;
+
+d = abs(x - p0);
+
+c0 = 2 .* (x>=-0.75 & x <= -0.25);
+c1 = 2 .* (x>=0.25 & x <= 0.75);
+c2 = 1.25 .* (x>=0.1 & x <= 0.9);
+
+peak_loss0 = sum(d.*c0) / length(x)
+peak_loss1 = sum(d.*c1) / length(x)
+peak_loss2 = sum(d.*c2) / length(x)
+
+createfigure(x, [c0;c1;c2], d, peak_loss0,peak_loss1, peak_loss2);

third_party/ALIKE/requirements.txt

@@ -0,0 +1,6 @@
+opencv-python~=4.5.1.48
+numpy~=1.19.5
+tqdm~=4.60.0
+torch~=1.8.0
+torchvision~=0.9.0
+thop~=0.0.31-2005241907

third_party/ALIKE/soft_detect.py

@@ -0,0 +1,234 @@
+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

third_party/ASpanFormer/.github/workflows/sync.yml

@@ -0,0 +1,39 @@
+name: Upstream Sync
+
+permissions:
+  contents: write
+
+on:
+  schedule:
+    - cron: "0 0 * * *" # every day
+  workflow_dispatch:
+
+jobs:
+  sync_latest_from_upstream:
+    name: Sync latest commits from upstream repo
+    runs-on: ubuntu-latest
+    if: ${{ github.event.repository.fork }}
+
+    steps:
+      # Step 1: run a standard checkout action
+      - name: Checkout target repo
+        uses: actions/checkout@v3
+
+      # Step 2: run the sync action
+      - name: Sync upstream changes
+        id: sync
+        uses: aormsby/Fork-Sync-With-Upstream-action@v3.4
+        with:
+          upstream_sync_repo: apple/ml-aspanformer
+          upstream_sync_branch: main
+          target_sync_branch: main
+          target_repo_token: ${{ secrets.GITHUB_TOKEN }} # automatically generated, no need to set
+
+          # Set test_mode true to run tests instead of the true action!!
+          test_mode: false
+
+      - name: Sync check
+        if: failure()
+        run: |
+          echo "::error::Due to insufficient permissions, synchronization failed (as expected). Please go to the repository homepage and manually perform [Sync fork]."
+          exit 1

third_party/ASpanFormer/.gitignore

@@ -0,0 +1,32 @@
+.vscode/
+__pycache__/
+*.pyc
+*.DS_Store
+*.swp
+*.pth
+tmp.*
+*/.ipynb_checkpoints/*
+
+logs/
+# weights/
+dump/
+demo/*.mp4
+demo/demo_images/
+src/loftr/utils/superglue.py
+demo/utils.py
+
+demo/*.jpg
+demo/*.png
+
+notebooks/QccDayNight.ipynb
+notebooks/westlake.ipynb
+assets/westlake
+assets/qcc_pairs.txt
+configs/.petrel*
+tools/draw_QccDayNights.py
+
+scripts/slurm/
+scripts/sbatch_submit.sh
+src/utils/client.py
+
+scannet_indices/

third_party/ASpanFormer/CODE_OF_CONDUCT.md

@@ -0,0 +1,71 @@
+# Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+  advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces, and it also applies when
+an individual is representing the project or its community in public spaces.
+Examples of representing a project or community include using an official
+project e-mail address, posting via an official social media account, or acting
+as an appointed representative at an online or offline event. Representation of
+a project may be further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the open source team at [opensource-conduct@group.apple.com](mailto:opensource-conduct@group.apple.com). All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org), version 1.4,
+available at [https://www.contributor-covenant.org/version/1/4/code-of-conduct.html](https://www.contributor-covenant.org/version/1/4/code-of-conduct.html)

third_party/ASpanFormer/CONTRIBUTING.md

@@ -0,0 +1,7 @@
+# Contribution Guide
+
+Thanks for your interest in contributing. This project was released to accompany a research paper for purposes of reproducability, and beyond its publication there are limited plans for future development of the repository.
+
+## Before you get started
+
+We ask that all community members read and observe our [Code of Conduct](CODE_OF_CONDUCT.md).

third_party/ASpanFormer/LICENSE

@@ -0,0 +1,9 @@
+Copyright (C) 2021, 2022 Apple Inc. All Rights Reserved. 
+
+IMPORTANT:  This Apple software is supplied to you by Apple Inc. ("Apple") in consideration of your agreement to the following terms, and your use, installation, modification or redistribution of this Apple software constitutes acceptance of these terms.  If you do not agree with these terms, please do not use, install, modify or redistribute this Apple software.
+
+In consideration of your agreement to abide by the following terms, and subject to these terms, Apple grants you a personal, non-commercial, non-exclusive license, under Apple's copyrights in this original Apple software (the "Apple Software"), to use, reproduce, modify and redistribute the Apple Software, with or without modifications, in source and/or binary forms for non-commercial purposes only; provided that if you redistribute the Apple Software in its entirety and without modifications, you must retain this notice and the following text and disclaimers in all such redistributions of the Apple Software. Neither the name, trademarks, service marks or logos of Apple Inc. may be used to endorse or promote products derived from the Apple Software without specific prior written permission from Apple.  Except as expressly stated in this notice, no other rights or licenses, express or implied, are granted by Apple herein, including but not limited to any patent rights that may be infringed by your derivative works or by other works in which the Apple Software may be incorporated.
+
+The Apple Software is provided by Apple on an "AS IS" basis.  APPLE MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND OPERATION ALONE OR IN COMBINATION WITH YOUR PRODUCTS.
+
+IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION, MODIFICATION AND/OR DISTRIBUTION OF THE APPLE SOFTWARE, HOWEVER CAUSED AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE), STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

third_party/ASpanFormer/README.md

@@ -0,0 +1,98 @@
+# Submodule used in [hloc](https://github.com/Vincentqyw/Hierarchical-Localization) toolbox
+
+# ASpanFormer Implementation
+
+![Framework](assets/teaser.png)
+
+This is a PyTorch implementation of ASpanFormer for ECCV'22 [paper](https://arxiv.org/abs/2208.14201), “ASpanFormer: Detector-Free Image Matching with Adaptive Span Transformer”, and can be used to reproduce the results in the paper.
+
+This work focuses on detector-free image matching. We propose a hierarchical attention framework for cross-view feature update, which adaptively adjusts attention span based on region-wise matchability.
+
+This repo contains training, evaluation and basic demo scripts used in our paper.
+
+A large part of the code base is borrowed from the [LoFTR Repository](https://github.com/zju3dv/LoFTR) under its own separate license, terms and conditions.  The authors of this software are not responsible for the contents of third-party websites.
+
+## Installation 
+```bash
+conda env create -f environment.yaml
+conda activate ASpanFormer
+```
+
+## Get started
+Download model weights from [here](https://drive.google.com/file/d/1eavM9dTkw9nbc-JqlVVfGPU5UvTTfc6k/view?usp=share_link)  
+
+Extract weights by
+```bash
+tar -xvf weights_aspanformer.tar
+```
+
+A demo to match one image pair is provided. To get a quick start, 
+
+```bash
+cd demo
+python demo.py
+```
+
+
+## Data Preparation
+Please follow the [training doc](docs/TRAINING.md) for data organization
+
+
+
+## Evaluation
+
+
+### 1. ScanNet Evaluation 
+```bash
+cd scripts/reproduce_test
+bash indoor.sh
+```
+Similar results as below should be obtained,
+```bash
+'auc@10': 0.46640095171012563,
+'auc@20': 0.6407042320049785,
+'auc@5': 0.26241231577189295,
+'prec@5e-04': 0.8827665604024288,
+'prec_flow@2e-03': 0.810938751342228
+```
+
+### 2. MegaDepth Evaluation
+ ```bash
+cd scripts/reproduce_test
+bash outdoor.sh
+```
+Similar results as below should be obtained,
+```bash
+'auc@10': 0.7184113573584142,
+'auc@20': 0.8333835724453831,
+'auc@5': 0.5567622479156181,
+'prec@5e-04': 0.9901741341790503,
+'prec_flow@2e-03': 0.7188964321862907
+```
+
+
+## Training
+
+### 1. ScanNet Training
+```bash
+cd scripts/reproduce_train
+bash indoor.sh
+```
+
+### 2. MegaDepth Training
+```bash
+cd scripts/reproduce_train
+bash outdoor.sh
+```
+      
+
+If you find this project useful, please cite:
+
+```
+@article{chen2022aspanformer,
+  title={ASpanFormer: Detector-Free Image Matching with Adaptive Span Transformer},
+  author={Chen, Hongkai and Luo, Zixin and Zhou, Lei and Tian, Yurun and Zhen, Mingmin and Fang, Tian and McKinnon, David and Tsin, Yanghai and Quan, Long},
+  journal={European Conference on Computer Vision (ECCV)},
+  year={2022}
+}
+```

third_party/ASpanFormer/configs/aspan/indoor/aspan_test.py

@@ -0,0 +1,11 @@
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent / "../../../"))
+from src.config.default import _CN as cfg
+
+cfg.ASPAN.MATCH_COARSE.MATCH_TYPE = "dual_softmax"
+
+cfg.ASPAN.MATCH_COARSE.BORDER_RM = 0
+cfg.ASPAN.COARSE.COARSEST_LEVEL = [15, 20]
+cfg.ASPAN.COARSE.TRAIN_RES = [480, 640]

third_party/ASpanFormer/configs/aspan/indoor/aspan_train.py

@@ -0,0 +1,12 @@
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent / "../../../"))
+from src.config.default import _CN as cfg
+
+cfg.ASPAN.COARSE.COARSEST_LEVEL = [15, 20]
+cfg.ASPAN.MATCH_COARSE.MATCH_TYPE = "dual_softmax"
+
+cfg.ASPAN.MATCH_COARSE.SPARSE_SPVS = False
+cfg.ASPAN.MATCH_COARSE.BORDER_RM = 0
+cfg.TRAINER.MSLR_MILESTONES = [3, 6, 9, 12, 17, 20, 23, 26, 29]

third_party/ASpanFormer/configs/aspan/outdoor/aspan_test.py

@@ -0,0 +1,22 @@
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent / "../../../"))
+from src.config.default import _CN as cfg
+
+cfg.ASPAN.COARSE.COARSEST_LEVEL = [36, 36]
+cfg.ASPAN.COARSE.TRAIN_RES = [832, 832]
+cfg.ASPAN.COARSE.TEST_RES = [1152, 1152]
+cfg.ASPAN.MATCH_COARSE.MATCH_TYPE = "dual_softmax"
+
+cfg.TRAINER.CANONICAL_LR = 8e-3
+cfg.TRAINER.WARMUP_STEP = 1875  # 3 epochs
+cfg.TRAINER.WARMUP_RATIO = 0.1
+cfg.TRAINER.MSLR_MILESTONES = [8, 12, 16, 20, 24]
+
+# pose estimation
+cfg.TRAINER.RANSAC_PIXEL_THR = 0.5
+
+cfg.TRAINER.OPTIMIZER = "adamw"
+cfg.TRAINER.ADAMW_DECAY = 0.1
+cfg.ASPAN.MATCH_COARSE.TRAIN_COARSE_PERCENT = 0.3

third_party/ASpanFormer/configs/aspan/outdoor/aspan_train.py

@@ -0,0 +1,21 @@
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent / "../../../"))
+from src.config.default import _CN as cfg
+
+cfg.ASPAN.COARSE.COARSEST_LEVEL = [26, 26]
+cfg.ASPAN.MATCH_COARSE.MATCH_TYPE = "dual_softmax"
+cfg.ASPAN.MATCH_COARSE.SPARSE_SPVS = False
+
+cfg.TRAINER.CANONICAL_LR = 8e-3
+cfg.TRAINER.WARMUP_STEP = 1875  # 3 epochs
+cfg.TRAINER.WARMUP_RATIO = 0.1
+cfg.TRAINER.MSLR_MILESTONES = [8, 12, 16, 20, 24]
+
+# pose estimation
+cfg.TRAINER.RANSAC_PIXEL_THR = 0.5
+
+cfg.TRAINER.OPTIMIZER = "adamw"
+cfg.TRAINER.ADAMW_DECAY = 0.1
+cfg.ASPAN.MATCH_COARSE.TRAIN_COARSE_PERCENT = 0.3

third_party/ASpanFormer/configs/data/base.py

@@ -0,0 +1,36 @@
+"""
+The data config will be the last one merged into the main config.
+Setups in data configs will override all existed setups!
+"""
+
+from yacs.config import CfgNode as CN
+
+_CN = CN()
+_CN.DATASET = CN()
+_CN.TRAINER = CN()
+
+# training data config
+_CN.DATASET.TRAIN_DATA_ROOT = None
+_CN.DATASET.TRAIN_POSE_ROOT = None
+_CN.DATASET.TRAIN_NPZ_ROOT = None
+_CN.DATASET.TRAIN_LIST_PATH = None
+_CN.DATASET.TRAIN_INTRINSIC_PATH = None
+# validation set config
+_CN.DATASET.VAL_DATA_ROOT = None
+_CN.DATASET.VAL_POSE_ROOT = None
+_CN.DATASET.VAL_NPZ_ROOT = None
+_CN.DATASET.VAL_LIST_PATH = None
+_CN.DATASET.VAL_INTRINSIC_PATH = None
+
+# testing data config
+_CN.DATASET.TEST_DATA_ROOT = None
+_CN.DATASET.TEST_POSE_ROOT = None
+_CN.DATASET.TEST_NPZ_ROOT = None
+_CN.DATASET.TEST_LIST_PATH = None
+_CN.DATASET.TEST_INTRINSIC_PATH = None
+
+# dataset config
+_CN.DATASET.MIN_OVERLAP_SCORE_TRAIN = 0.4
+_CN.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0  # for both test and val
+
+cfg = _CN

third_party/ASpanFormer/configs/data/debug/.gitignore

@@ -0,0 +1,3 @@
+*
+*/
+!.gitignore

third_party/ASpanFormer/configs/data/megadepth_test_1500.py

@@ -0,0 +1,13 @@
+from configs.data.base import cfg
+
+TEST_BASE_PATH = "assets/megadepth_test_1500_scene_info"
+
+cfg.DATASET.TEST_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.TEST_DATA_ROOT = "data/megadepth/test"
+cfg.DATASET.TEST_NPZ_ROOT = f"{TEST_BASE_PATH}"
+cfg.DATASET.TEST_LIST_PATH = f"{TEST_BASE_PATH}/megadepth_test_1500.txt"
+
+cfg.DATASET.MGDPT_IMG_RESIZE = 1152
+cfg.DATASET.MGDPT_IMG_PAD = True
+cfg.DATASET.MGDPT_DF = 8
+cfg.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0

third_party/ASpanFormer/configs/data/megadepth_trainval_832.py

@@ -0,0 +1,26 @@
+from configs.data.base import cfg
+
+
+TRAIN_BASE_PATH = "data/megadepth/index"
+cfg.DATASET.TRAINVAL_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.TRAIN_DATA_ROOT = "data/megadepth/train"
+cfg.DATASET.TRAIN_NPZ_ROOT = f"{TRAIN_BASE_PATH}/scene_info_0.1_0.7"
+cfg.DATASET.TRAIN_LIST_PATH = f"{TRAIN_BASE_PATH}/trainvaltest_list/train_list.txt"
+cfg.DATASET.MIN_OVERLAP_SCORE_TRAIN = 0.0
+
+TEST_BASE_PATH = "data/megadepth/index"
+cfg.DATASET.TEST_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.VAL_DATA_ROOT = cfg.DATASET.TEST_DATA_ROOT = "data/megadepth/test"
+cfg.DATASET.VAL_NPZ_ROOT = (
+    cfg.DATASET.TEST_NPZ_ROOT
+) = f"{TEST_BASE_PATH}/scene_info_val_1500"
+cfg.DATASET.VAL_LIST_PATH = (
+    cfg.DATASET.TEST_LIST_PATH
+) = f"{TEST_BASE_PATH}/trainvaltest_list/val_list.txt"
+cfg.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0  # for both test and val
+
+# 368 scenes in total for MegaDepth
+# (with difficulty balanced (further split each scene to 3 sub-scenes))
+cfg.TRAINER.N_SAMPLES_PER_SUBSET = 100
+
+cfg.DATASET.MGDPT_IMG_RESIZE = 832  # for training on 32GB meme GPUs

third_party/ASpanFormer/configs/data/scannet_test_1500.py

@@ -0,0 +1,11 @@
+from configs.data.base import cfg
+
+TEST_BASE_PATH = "assets/scannet_test_1500"
+
+cfg.DATASET.TEST_DATA_SOURCE = "ScanNet"
+cfg.DATASET.TEST_DATA_ROOT = "data/scannet/test"
+cfg.DATASET.TEST_NPZ_ROOT = f"{TEST_BASE_PATH}"
+cfg.DATASET.TEST_LIST_PATH = f"{TEST_BASE_PATH}/scannet_test.txt"
+cfg.DATASET.TEST_INTRINSIC_PATH = f"{TEST_BASE_PATH}/intrinsics.npz"
+
+cfg.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0

third_party/ASpanFormer/configs/data/scannet_trainval.py

@@ -0,0 +1,21 @@
+from configs.data.base import cfg
+
+
+TRAIN_BASE_PATH = "data/scannet/index"
+cfg.DATASET.TRAINVAL_DATA_SOURCE = "ScanNet"
+cfg.DATASET.TRAIN_DATA_ROOT = "data/scannet/train"
+cfg.DATASET.TRAIN_NPZ_ROOT = f"{TRAIN_BASE_PATH}/scene_data/train"
+cfg.DATASET.TRAIN_LIST_PATH = f"{TRAIN_BASE_PATH}/scene_data/train_list/scannet_all.txt"
+cfg.DATASET.TRAIN_INTRINSIC_PATH = f"{TRAIN_BASE_PATH}/intrinsics.npz"
+
+TEST_BASE_PATH = "assets/scannet_test_1500"
+cfg.DATASET.TEST_DATA_SOURCE = "ScanNet"
+cfg.DATASET.VAL_DATA_ROOT = cfg.DATASET.TEST_DATA_ROOT = "data/scannet/test"
+cfg.DATASET.VAL_NPZ_ROOT = cfg.DATASET.TEST_NPZ_ROOT = TEST_BASE_PATH
+cfg.DATASET.VAL_LIST_PATH = (
+    cfg.DATASET.TEST_LIST_PATH
+) = f"{TEST_BASE_PATH}/scannet_test.txt"
+cfg.DATASET.VAL_INTRINSIC_PATH = (
+    cfg.DATASET.TEST_INTRINSIC_PATH
+) = f"{TEST_BASE_PATH}/intrinsics.npz"
+cfg.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0  # for both test and val

third_party/ASpanFormer/data/megadepth/index/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

third_party/ASpanFormer/data/megadepth/test/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

third_party/ASpanFormer/data/megadepth/train/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

third_party/ASpanFormer/data/scannet/index/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

third_party/ASpanFormer/data/scannet/test/.gitignore

@@ -0,0 +1,3 @@
+*
+*/
+!.gitignore

third_party/ASpanFormer/data/scannet/train/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

third_party/ASpanFormer/demo/demo.py

@@ -0,0 +1,91 @@
+import os
+import sys
+
+ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, ROOT_DIR)
+
+from src.ASpanFormer.aspanformer import ASpanFormer
+from src.config.default import get_cfg_defaults
+from src.utils.misc import lower_config
+import demo_utils
+
+import cv2
+import torch
+import numpy as np
+
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--config_path",
+    type=str,
+    default="../configs/aspan/outdoor/aspan_test.py",
+    help="path for config file.",
+)
+parser.add_argument(
+    "--img0_path",
+    type=str,
+    default="../assets/phototourism_sample_images/piazza_san_marco_06795901_3725050516.jpg",
+    help="path for image0.",
+)
+parser.add_argument(
+    "--img1_path",
+    type=str,
+    default="../assets/phototourism_sample_images/piazza_san_marco_15148634_5228701572.jpg",
+    help="path for image1.",
+)
+parser.add_argument(
+    "--weights_path",
+    type=str,
+    default="../weights/outdoor.ckpt",
+    help="path for model weights.",
+)
+parser.add_argument(
+    "--long_dim0", type=int, default=1024, help="resize for longest dim of image0."
+)
+parser.add_argument(
+    "--long_dim1", type=int, default=1024, help="resize for longest dim of image1."
+)
+
+args = parser.parse_args()
+
+
+if __name__ == "__main__":
+    config = get_cfg_defaults()
+    config.merge_from_file(args.config_path)
+    _config = lower_config(config)
+    matcher = ASpanFormer(config=_config["aspan"])
+    state_dict = torch.load(args.weights_path, map_location="cpu")["state_dict"]
+    matcher.load_state_dict(state_dict, strict=False)
+    matcher.cuda(), matcher.eval()
+
+    img0, img1 = cv2.imread(args.img0_path), cv2.imread(args.img1_path)
+    img0_g, img1_g = cv2.imread(args.img0_path, 0), cv2.imread(args.img1_path, 0)
+    img0, img1 = demo_utils.resize(img0, args.long_dim0), demo_utils.resize(
+        img1, args.long_dim1
+    )
+    img0_g, img1_g = demo_utils.resize(img0_g, args.long_dim0), demo_utils.resize(
+        img1_g, args.long_dim1
+    )
+    data = {
+        "image0": torch.from_numpy(img0_g / 255.0)[None, None].cuda().float(),
+        "image1": torch.from_numpy(img1_g / 255.0)[None, None].cuda().float(),
+    }
+    with torch.no_grad():
+        matcher(data, online_resize=True)
+        corr0, corr1 = data["mkpts0_f"].cpu().numpy(), data["mkpts1_f"].cpu().numpy()
+
+    F_hat, mask_F = cv2.findFundamentalMat(
+        corr0, corr1, method=cv2.FM_RANSAC, ransacReprojThreshold=1
+    )
+    if mask_F is not None:
+        mask_F = mask_F[:, 0].astype(bool)
+    else:
+        mask_F = np.zeros_like(corr0[:, 0]).astype(bool)
+
+    # visualize match
+    display = demo_utils.draw_match(img0, img1, corr0, corr1)
+    display_ransac = demo_utils.draw_match(img0, img1, corr0[mask_F], corr1[mask_F])
+    cv2.imwrite("match.png", display)
+    cv2.imwrite("match_ransac.png", display_ransac)
+    print(len(corr1), len(corr1[mask_F]))

third_party/ASpanFormer/demo/demo_utils.py

@@ -0,0 +1,88 @@
+import cv2
+import numpy as np
+
+
+def resize(image, long_dim):
+    h, w = image.shape[0], image.shape[1]
+    image = cv2.resize(
+        image, (int(w * long_dim / max(h, w)), int(h * long_dim / max(h, w)))
+    )
+    return image
+
+
+def draw_points(img, points, color=(0, 255, 0), radius=3):
+    dp = [(int(points[i, 0]), int(points[i, 1])) for i in range(points.shape[0])]
+    for i in range(points.shape[0]):
+        cv2.circle(img, dp[i], radius=radius, color=color)
+    return img
+
+
+def draw_match(
+    img1,
+    img2,
+    corr1,
+    corr2,
+    inlier=[True],
+    color=None,
+    radius1=1,
+    radius2=1,
+    resize=None,
+):
+    if resize is not None:
+        scale1, scale2 = [img1.shape[1] / resize[0], img1.shape[0] / resize[1]], [
+            img2.shape[1] / resize[0],
+            img2.shape[0] / resize[1],
+        ]
+        img1, img2 = cv2.resize(img1, resize, interpolation=cv2.INTER_AREA), cv2.resize(
+            img2, resize, interpolation=cv2.INTER_AREA
+        )
+        corr1, corr2 = (
+            corr1 / np.asarray(scale1)[np.newaxis],
+            corr2 / np.asarray(scale2)[np.newaxis],
+        )
+    corr1_key = [
+        cv2.KeyPoint(corr1[i, 0], corr1[i, 1], radius1) for i in range(corr1.shape[0])
+    ]
+    corr2_key = [
+        cv2.KeyPoint(corr2[i, 0], corr2[i, 1], radius2) for i in range(corr2.shape[0])
+    ]
+
+    assert len(corr1) == len(corr2)
+
+    draw_matches = [cv2.DMatch(i, i, 0) for i in range(len(corr1))]
+    if color is None:
+        color = [(0, 255, 0) if cur_inlier else (0, 0, 255) for cur_inlier in inlier]
+    if len(color) == 1:
+        display = cv2.drawMatches(
+            img1,
+            corr1_key,
+            img2,
+            corr2_key,
+            draw_matches,
+            None,
+            matchColor=color[0],
+            singlePointColor=color[0],
+            flags=4,
+        )
+    else:
+        height, width = max(img1.shape[0], img2.shape[0]), img1.shape[1] + img2.shape[1]
+        display = np.zeros([height, width, 3], np.uint8)
+        display[: img1.shape[0], : img1.shape[1]] = img1
+        display[: img2.shape[0], img1.shape[1] :] = img2
+        for i in range(len(corr1)):
+            left_x, left_y, right_x, right_y = (
+                int(corr1[i][0]),
+                int(corr1[i][1]),
+                int(corr2[i][0] + img1.shape[1]),
+                int(corr2[i][1]),
+            )
+            cur_color = (int(color[i][0]), int(color[i][1]), int(color[i][2]))
+            cv2.line(
+                display,
+                (left_x, left_y),
+                (right_x, right_y),
+                cur_color,
+                1,
+                lineType=cv2.LINE_AA,
+            )
+    return display

third_party/ASpanFormer/docs/TRAINING.md

@@ -0,0 +1,72 @@
+
+# Traininig ASpanFormer
+
+## Dataset setup
+Generally, two parts of data are needed for training ASpanFormer, the original dataset, i.e., ScanNet and MegaDepth, and the offline generated dataset indices. The dataset indices store scenes, image pairs, and other metadata within each dataset used for training/validation/testing. For the MegaDepth dataset, the relative poses between images used for training are directly cached in the indexing files. However, the relative poses of ScanNet image pairs are not stored due to the enormous resulting file size.
+
+### Download datasets
+#### MegaDepth
+We use depth maps provided in the [original MegaDepth dataset](https://www.cs.cornell.edu/projects/megadepth/) as well as undistorted images, corresponding camera intrinsics and extrinsics preprocessed by [D2-Net](https://github.com/mihaidusmanu/d2-net#downloading-and-preprocessing-the-megadepth-dataset). You can download them separately from the following links. 
+- [MegaDepth undistorted images and processed depths](https://www.cs.cornell.edu/projects/megadepth/dataset/Megadepth_v1/MegaDepth_v1.tar.gz)
+    - Note that we only use depth maps.
+    - Path of the download data will be referreed to as `/path/to/megadepth`
+- [D2-Net preprocessed images](https://drive.google.com/drive/folders/1hxpOsqOZefdrba_BqnW490XpNX_LgXPB)
+    - Images are undistorted manually in D2-Net since the undistorted images from MegaDepth do not come with corresponding intrinsics.
+    - Path of the download data will be referreed to as `/path/to/megadepth_d2net`
+
+#### ScanNet
+Please set up the ScanNet dataset following [the official guide](https://github.com/ScanNet/ScanNet#scannet-data)
+> NOTE: We use the [python exported data](https://github.com/ScanNet/ScanNet/tree/master/SensReader/python),
+instead of the [c++ exported one](https://github.com/ScanNet/ScanNet/tree/master/SensReader/c%2B%2B).
+
+### Download the dataset indices
+
+You can download the required dataset indices from the [following link](https://drive.google.com/drive/folders/1DOcOPZb3-5cWxLqn256AhwUVjBPifhuf).
+After downloading, unzip the required files.
+```shell
+unzip downloaded-file.zip
+
+# extract dataset indices
+tar xf train-data/megadepth_indices.tar
+tar xf train-data/scannet_indices.tar
+
+# extract testing data (optional)
+tar xf testdata/megadepth_test_1500.tar
+tar xf testdata/scannet_test_1500.tar
+```
+
+### Build the dataset symlinks
+
+We symlink the datasets to the `data` directory under the main ASpanFormer project directory.
+
+```shell
+# scannet
+# -- # train and test dataset
+ln -s /path/to/scannet_train/* /path/to/ASpanFormer/data/scannet/train
+ln -s /path/to/scannet_test/* /path/to/ASpanFormer/data/scannet/test
+# -- # dataset indices
+ln -s /path/to/scannet_indices/* /path/to/ASpanFormer/data/scannet/index
+
+# megadepth
+# -- # train and test dataset (train and test share the same dataset)
+ln -sv /path/to/megadepth/phoenix /path/to/megadepth_d2net/Undistorted_SfM /path/to/ASpanFormer/data/megadepth/train
+ln -sv /path/to/megadepth/phoenix /path/to/megadepth_d2net/Undistorted_SfM /path/to/ASpanFormer/data/megadepth/test
+# -- # dataset indices
+ln -s /path/to/megadepth_indices/* /path/to/ASpanFormer/data/megadepth/index
+```
+
+
+## Training
+We provide training scripts of ScanNet and MegaDepth. The results in the ASpanFormer paper can be reproduced with 8 v100 GPUs. For a different setup, we scale the learning rate and its warm-up linearly, but the final evaluation results might vary due to the different batch size & learning rate used. Thus the reproduction of results in our paper is not guaranteed.
+
+
+### Training on ScanNet
+``` shell
+scripts/reproduce_train/indoor.sh
+```
+
+
+### Training on MegaDepth
+``` shell
+scripts/reproduce_train/outdoor.sh
+```

third_party/ASpanFormer/environment.yaml

@@ -0,0 +1,12 @@
+name: ASpanFormer
+channels:
+  - pytorch
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.8
+  - cudatoolkit=10.2
+  - pytorch=1.8.1
+  - pip
+  - pip:
+      - -r requirements.txt

third_party/ASpanFormer/requirements.txt

@@ -0,0 +1,18 @@
+#opencv_python==4.4.0.46
+albumentations==0.5.1 --no-binary=imgaug,albumentations
+ray>=1.0.1
+einops==0.3.0
+kornia==0.4.1
+loguru==0.5.3
+yacs>=0.1.8
+tqdm
+autopep8
+pylint
+ipython
+jupyterlab
+matplotlib
+h5py
+pytorch-lightning==1.3.5
+loguru
+joblib>=1.0.1
+torchmetrics==0.4

third_party/ASpanFormer/scripts/reproduce_test/indoor.sh

@@ -0,0 +1,31 @@
+#!/bin/bash -l
+# a indoor_ds model with the pos_enc impl bug fixed.
+
+SCRIPTPATH=$(dirname $(readlink -f "$0"))
+PROJECT_DIR="${SCRIPTPATH}/../../"
+
+# conda activate loftr
+export PYTHONPATH=$PROJECT_DIR:$PYTHONPATH
+cd $PROJECT_DIR
+
+data_cfg_path="configs/data/scannet_test_1500.py"
+main_cfg_path="configs/aspan/indoor/aspan_test.py"
+ckpt_path='weights/indoor.ckpt'
+dump_dir="dump/indoor_dump"
+profiler_name="inference"
+n_nodes=1  # mannually keep this the same with --nodes
+n_gpus_per_node=-1
+torch_num_workers=4
+batch_size=1  # per gpu
+
+python -u ./test.py \
+    ${data_cfg_path} \
+    ${main_cfg_path} \
+    --ckpt_path=${ckpt_path} \
+    --dump_dir=${dump_dir} \
+    --gpus=${n_gpus_per_node} --num_nodes=${n_nodes} --accelerator="ddp" \
+    --batch_size=${batch_size} --num_workers=${torch_num_workers}\
+    --profiler_name=${profiler_name} \
+    --benchmark \
+    --mode integrated
+