init: files

.gitignore

@@ -0,0 +1,28 @@
+build/
+# lib
+bin/
+cmake_modules/
+cmake-build-debug/
+.idea/
+.vscode/
+*.pyc
+flagged
+.ipynb_checkpoints
+__pycache__
+Untitled*
+experiments
+third_party/REKD
+hloc/matchers/dedode.py
+gradio_cached_examples
+*.mp4
+hloc/matchers/quadtree.py
+third_party/QuadTreeAttention
+desktop.ini
+*.egg-info
+output.pkl
+log.txt
+experiments*
+gen_example.py
+datasets/lines/terrace0.JPG
+datasets/lines/terrace1.JPG
+datasets/South-Building*

Dockerfile

@@ -0,0 +1,27 @@
+# Use an official conda-based Python image as a parent image
+FROM pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime
+LABEL maintainer vincentqyw
+ARG PYTHON_VERSION=3.10.10
+
+# Set the working directory to /code
+WORKDIR /code
+
+# Install Git and Git LFS
+RUN apt-get update && apt-get install -y git-lfs
+RUN git lfs install
+
+# Clone the Git repository
+RUN git clone https://huggingface.co/spaces/Realcat/image-matching-webui /code
+
+RUN conda create -n imw python=${PYTHON_VERSION}
+RUN echo "source activate imw" > ~/.bashrc
+ENV PATH /opt/conda/envs/imw/bin:$PATH
+
+# Make RUN commands use the new environment
+SHELL ["conda", "run", "-n", "imw", "/bin/bash", "-c"]
+RUN pip install --upgrade pip
+RUN pip install -r requirements.txt
+RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
+
+# Export port
+EXPOSE 7860

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,12 +1,155 @@
 ---
 title: Imatchui
-emoji: 🐨
+emoji: 🤗
 colorFrom: red
-colorTo: indigo
+colorTo: yellow
 sdk: gradio
 sdk_version: 5.4.0
 app_file: app.py
-pinned: false
+pinned: true
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+[![Contributors][contributors-shield]][contributors-url]
+[![Forks][forks-shield]][forks-url]
+[![Stargazers][stars-shield]][stars-url]
+[![Issues][issues-shield]][issues-url]
+
+<p align="center">
+  <h1 align="center"><br><ins>Image Matching WebUI</ins><br>Identify matching points between two images</h1> 
+</p>
+
+## Description
+
+This simple tool efficiently matches image pairs using multiple famous image matching algorithms. The tool features a Graphical User Interface (GUI) designed using [gradio](https://gradio.app/). You can effortlessly select two images and a matching algorithm and obtain a precise matching result.
+**Note**: the images source can be either local images or webcam images.
+
+Try it on <a href='https://huggingface.co/spaces/Realcat/image-matching-webui'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue'></a> 
+<a target="_blank" href="https://lightning.ai/realcat/studios/image-matching-webui">
+  <img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/>
+</a>
+
+Here is a demo of the tool:
+
+![demo](assets/demo.gif)
+
+The tool currently supports various popular image matching algorithms, namely:
+- [x] [EfficientLoFTR](https://github.com/zju3dv/EfficientLoFTR), CVPR 2024
+- [x] [MASt3R](https://github.com/naver/mast3r), CVPR 2024
+- [x] [DUSt3R](https://github.com/naver/dust3r), CVPR 2024
+- [x] [OmniGlue](https://github.com/Vincentqyw/omniglue-onnx), CVPR 2024
+- [x] [XFeat](https://github.com/verlab/accelerated_features), CVPR 2024
+- [x] [RoMa](https://github.com/Vincentqyw/RoMa), CVPR 2024
+- [x] [DeDoDe](https://github.com/Parskatt/DeDoDe), 3DV 2024
+- [ ] [Mickey](https://github.com/nianticlabs/mickey), CVPR 2024
+- [x] [GIM](https://github.com/xuelunshen/gim), ICLR 2024
+- [ ] [DUSt3R](https://github.com/naver/dust3r), arXiv 2023
+- [x] [LightGlue](https://github.com/cvg/LightGlue), ICCV 2023
+- [x] [DarkFeat](https://github.com/THU-LYJ-Lab/DarkFeat), AAAI 2023
+- [x] [SFD2](https://github.com/feixue94/sfd2), CVPR 2023
+- [x] [IMP](https://github.com/feixue94/imp-release), CVPR 2023
+- [ ] [ASTR](https://github.com/ASTR2023/ASTR), CVPR 2023
+- [ ] [SEM](https://github.com/SEM2023/SEM), CVPR 2023
+- [ ] [DeepLSD](https://github.com/cvg/DeepLSD), CVPR 2023
+- [x] [GlueStick](https://github.com/cvg/GlueStick), ICCV 2023
+- [ ] [ConvMatch](https://github.com/SuhZhang/ConvMatch), AAAI 2023
+- [x] [LoFTR](https://github.com/zju3dv/LoFTR), CVPR 2021
+- [x] [SOLD2](https://github.com/cvg/SOLD2), CVPR 2021
+- [ ] [LineTR](https://github.com/yosungho/LineTR), RA-L 2021
+- [x] [DKM](https://github.com/Parskatt/DKM), CVPR 2023
+- [ ] [NCMNet](https://github.com/xinliu29/NCMNet), CVPR 2023
+- [x] [TopicFM](https://github.com/Vincentqyw/TopicFM), AAAI 2023
+- [x] [AspanFormer](https://github.com/Vincentqyw/ml-aspanformer), ECCV 2022
+- [x] [LANet](https://github.com/wangch-g/lanet), ACCV 2022
+- [ ] [LISRD](https://github.com/rpautrat/LISRD), ECCV 2022
+- [ ] [REKD](https://github.com/bluedream1121/REKD), CVPR 2022
+- [x] [CoTR](https://github.com/ubc-vision/COTR), ICCV 2021
+- [x] [ALIKE](https://github.com/Shiaoming/ALIKE), TMM 2022
+- [x] [RoRD](https://github.com/UditSinghParihar/RoRD), IROS 2021
+- [x] [SGMNet](https://github.com/vdvchen/SGMNet), ICCV 2021
+- [x] [SuperPoint](https://github.com/magicleap/SuperPointPretrainedNetwork), CVPRW 2018
+- [x] [SuperGlue](https://github.com/magicleap/SuperGluePretrainedNetwork), CVPR 2020
+- [x] [D2Net](https://github.com/Vincentqyw/d2-net), CVPR 2019
+- [x] [R2D2](https://github.com/naver/r2d2), NeurIPS 2019
+- [x] [DISK](https://github.com/cvlab-epfl/disk), NeurIPS 2020
+- [ ] [Key.Net](https://github.com/axelBarroso/Key.Net), ICCV 2019
+- [ ] [OANet](https://github.com/zjhthu/OANet), ICCV 2019
+- [x] [SOSNet](https://github.com/scape-research/SOSNet), CVPR 2019
+- [x] [HardNet](https://github.com/DagnyT/hardnet), NeurIPS 2017
+- [x] [SIFT](https://docs.opencv.org/4.x/da/df5/tutorial_py_sift_intro.html), IJCV 2004
+
+## How to use
+
+### HuggingFace / Lightning AI
+
+Just try it on <a href='https://huggingface.co/spaces/Realcat/image-matching-webui'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue'></a> 
+<a target="_blank" href="https://lightning.ai/realcat/studios/image-matching-webui">
+  <img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/>
+</a>
+
+or deploy it locally following the instructions below.
+
+### Requirements
+``` bash
+git clone --recursive https://github.com/Vincentqyw/image-matching-webui.git
+cd image-matching-webui
+conda env create -f environment.yaml
+conda activate imw
+```
+
+or using [docker](https://hub.docker.com/r/vincentqin/image-matching-webui):
+
+``` bash
+docker pull vincentqin/image-matching-webui:latest
+docker run -it -p 7860:7860 vincentqin/image-matching-webui:latest python app.py --server_name "0.0.0.0" --server_port=7860
+```
+ 
+### Run demo
+``` bash
+python3 ./app.py
+```
+then open http://localhost:7860 in your browser.
+
+![](assets/gui.jpg)
+
+### Add your own feature / matcher
+
+I provide an example to add local feature in [hloc/extractors/example.py](hloc/extractors/example.py). Then add feature settings in `confs` in file [hloc/extract_features.py](hloc/extract_features.py). Last step is adding some settings to `model_zoo` in file [ui/config.yaml](ui/config.yaml).
+
+## Contributions welcome!
+
+External contributions are very much welcome. Please follow the [PEP8 style guidelines](https://www.python.org/dev/peps/pep-0008/) using a linter like flake8 (reformat using command `python -m black .`). This is a non-exhaustive list of features that might be valuable additions:
+
+- [x] add webcam support
+- [x] add [line feature matching](https://github.com/Vincentqyw/LineSegmentsDetection) algorithms
+- [x] example to add a new feature extractor / matcher
+- [x] ransac to filter outliers
+- [ ] add [rotation images](https://github.com/pidahbus/deep-image-orientation-angle-detection) options before matching 
+- [ ] support export matches to colmap ([#issue 6](https://github.com/Vincentqyw/image-matching-webui/issues/6))
+- [ ] add config file to set default parameters
+- [ ] dynamically load models and reduce GPU overload
+
+Adding local features / matchers as submodules is very easy. For example, to add the [GlueStick](https://github.com/cvg/GlueStick): 
+
+``` bash
+git submodule add https://github.com/cvg/GlueStick.git third_party/GlueStick
+```
+
+If remote submodule repositories are updated, don't forget to pull submodules with `git submodule update --remote`, if you only want to update one submodule, use `git submodule update --remote third_party/GlueStick`.
+
+## Resources
+- [Image Matching: Local Features & Beyond](https://image-matching-workshop.github.io)
+- [Long-term Visual Localization](https://www.visuallocalization.net)
+
+## Acknowledgement
+
+This code is built based on [Hierarchical-Localization](https://github.com/cvg/Hierarchical-Localization). We express our gratitude to the authors for their valuable source code.
+
+[contributors-shield]: https://img.shields.io/github/contributors/Vincentqyw/image-matching-webui.svg?style=for-the-badge
+[contributors-url]: https://github.com/Vincentqyw/image-matching-webui/graphs/contributors
+[forks-shield]: https://img.shields.io/github/forks/Vincentqyw/image-matching-webui.svg?style=for-the-badge
+[forks-url]: https://github.com/Vincentqyw/image-matching-webui/network/members
+[stars-shield]: https://img.shields.io/github/stars/Vincentqyw/image-matching-webui.svg?style=for-the-badge
+[stars-url]: https://github.com/Vincentqyw/image-matching-webui/stargazers
+[issues-shield]: https://img.shields.io/github/issues/Vincentqyw/image-matching-webui.svg?style=for-the-badge
+[issues-url]: https://github.com/Vincentqyw/image-matching-webui/issues

api/client.py

@@ -0,0 +1,225 @@
+import argparse
+import base64
+import os
+import pickle
+import time
+from typing import Dict, List
+
+import cv2
+import numpy as np
+import requests
+
+ENDPOINT = "http://127.0.0.1:8001"
+if "REMOTE_URL_RAILWAY" in os.environ:
+    ENDPOINT = os.environ["REMOTE_URL_RAILWAY"]
+
+print(f"API ENDPOINT: {ENDPOINT}")
+
+API_VERSION = f"{ENDPOINT}/version"
+API_URL_MATCH = f"{ENDPOINT}/v1/match"
+API_URL_EXTRACT = f"{ENDPOINT}/v1/extract"
+
+
+def read_image(path: str) -> str:
+    """
+    Read an image from a file, encode it as a JPEG and then as a base64 string.
+
+    Args:
+        path (str): The path to the image to read.
+
+    Returns:
+        str: The base64 encoded image.
+    """
+    # Read the image from the file
+    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+
+    # Encode the image as a png, NO COMPRESSION!!!
+    retval, buffer = cv2.imencode(".png", img)
+
+    # Encode the JPEG as a base64 string
+    b64img = base64.b64encode(buffer).decode("utf-8")
+
+    return b64img
+
+
+def do_api_requests(url=API_URL_EXTRACT, **kwargs):
+    """
+    Helper function to send an API request to the image matching service.
+
+    Args:
+        url (str): The URL of the API endpoint to use. Defaults to the
+            feature extraction endpoint.
+        **kwargs: Additional keyword arguments to pass to the API.
+
+    Returns:
+        List[Dict[str, np.ndarray]]: A list of dictionaries containing the
+            extracted features. The keys are "keypoints", "descriptors", and
+            "scores", and the values are ndarrays of shape (N, 2), (N, ?),
+            and (N,), respectively.
+    """
+    # Set up the request body
+    reqbody = {
+        # List of image data base64 encoded
+        "data": [],
+        # List of maximum number of keypoints to extract from each image
+        "max_keypoints": [100, 100],
+        # List of timestamps for each image (not used?)
+        "timestamps": ["0", "1"],
+        # Whether to convert the images to grayscale
+        "grayscale": 0,
+        # List of image height and width
+        "image_hw": [[640, 480], [320, 240]],
+        # Type of feature to extract
+        "feature_type": 0,
+        # List of rotation angles for each image
+        "rotates": [0.0, 0.0],
+        # List of scale factors for each image
+        "scales": [1.0, 1.0],
+        # List of reference points for each image (not used)
+        "reference_points": [[640, 480], [320, 240]],
+        # Whether to binarize the descriptors
+        "binarize": True,
+    }
+    # Update the request body with the additional keyword arguments
+    reqbody.update(kwargs)
+    try:
+        # Send the request
+        r = requests.post(url, json=reqbody)
+        if r.status_code == 200:
+            # Return the response
+            return r.json()
+        else:
+            # Print an error message if the response code is not 200
+            print(f"Error: Response code {r.status_code} - {r.text}")
+    except Exception as e:
+        # Print an error message if an exception occurs
+        print(f"An error occurred: {e}")
+
+
+def send_request_match(path0: str, path1: str) -> Dict[str, np.ndarray]:
+    """
+    Send a request to the API to generate a match between two images.
+
+    Args:
+        path0 (str): The path to the first image.
+        path1 (str): The path to the second image.
+
+    Returns:
+        Dict[str, np.ndarray]: A dictionary containing the generated matches.
+            The keys are "keypoints0", "keypoints1", "matches0", and "matches1",
+            and the values are ndarrays of shape (N, 2), (N, 2), (N, 2), and
+            (N, 2), respectively.
+    """
+    files = {"image0": open(path0, "rb"), "image1": open(path1, "rb")}
+    try:
+        # TODO: replace files with post json
+        response = requests.post(API_URL_MATCH, files=files)
+        pred = {}
+        if response.status_code == 200:
+            pred = response.json()
+            for key in list(pred.keys()):
+                pred[key] = np.array(pred[key])
+        else:
+            print(
+                f"Error: Response code {response.status_code} - {response.text}"
+            )
+    finally:
+        files["image0"].close()
+        files["image1"].close()
+    return pred
+
+
+def send_request_extract(
+    input_images: str, viz: bool = False
+) -> List[Dict[str, np.ndarray]]:
+    """
+    Send a request to the API to extract features from an image.
+
+    Args:
+        input_images (str): The path to the image.
+
+    Returns:
+        List[Dict[str, np.ndarray]]: A list of dictionaries containing the
+            extracted features. The keys are "keypoints", "descriptors", and
+            "scores", and the values are ndarrays of shape (N, 2), (N, 128),
+            and (N,), respectively.
+    """
+    image_data = read_image(input_images)
+    inputs = {
+        "data": [image_data],
+    }
+    response = do_api_requests(
+        url=API_URL_EXTRACT,
+        **inputs,
+    )
+    print("Keypoints detected: {}".format(len(response[0]["keypoints"])))
+
+    # draw matching, debug only
+    if viz:
+        from hloc.utils.viz import plot_keypoints
+        from ui.viz import fig2im, plot_images
+
+        kpts = np.array(response[0]["keypoints_orig"])
+        if "image_orig" in response[0].keys():
+            img_orig = np.array(["image_orig"])
+
+            output_keypoints = plot_images([img_orig], titles="titles", dpi=300)
+            plot_keypoints([kpts])
+            output_keypoints = fig2im(output_keypoints)
+            cv2.imwrite(
+                "demo_match.jpg",
+                output_keypoints[:, :, ::-1].copy(),  # RGB -> BGR
+            )
+    return response
+
+
+def get_api_version():
+    try:
+        response = requests.get(API_VERSION).json()
+        print("API VERSION: {}".format(response["version"]))
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Send text to stable audio server and receive generated audio."
+    )
+    parser.add_argument(
+        "--image0",
+        required=False,
+        help="Path for the file's melody",
+        default="datasets/sacre_coeur/mapping_rot/02928139_3448003521_rot45.jpg",
+    )
+    parser.add_argument(
+        "--image1",
+        required=False,
+        help="Path for the file's melody",
+        default="datasets/sacre_coeur/mapping_rot/02928139_3448003521_rot90.jpg",
+    )
+    args = parser.parse_args()
+
+    # get api version
+    get_api_version()
+
+    # request match
+    # for i in range(10):
+    #     t1 = time.time()
+    #     preds = send_request_match(args.image0, args.image1)
+    #     t2 = time.time()
+    #     print(
+    #         "Time cost1: {} seconds, matched: {}".format(
+    #             (t2 - t1), len(preds["mmkeypoints0_orig"])
+    #         )
+    #     )
+
+    # request extract
+    for i in range(10):
+        t1 = time.time()
+        preds = send_request_extract(args.image0)
+        t2 = time.time()
+        print(f"Time cost2: {(t2 - t1)} seconds")
+
+    # dump preds
+    with open("preds.pkl", "wb") as f:
+        pickle.dump(preds, f)

api/server.py

@@ -0,0 +1,499 @@
+# server.py
+import base64
+import io
+import sys
+import warnings
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI, File, UploadFile
+from fastapi.exceptions import HTTPException
+from fastapi.responses import JSONResponse
+from PIL import Image
+
+sys.path.append(str(Path(__file__).parents[1]))
+
+from api.types import ImagesInput
+from hloc import DEVICE, extract_features, logger, match_dense, match_features
+from hloc.utils.viz import add_text, plot_keypoints
+from ui import get_version
+from ui.utils import filter_matches, get_feature_model, get_model
+from ui.viz import display_matches, fig2im, plot_images
+
+warnings.simplefilter("ignore")
+
+
+def decode_base64_to_image(encoding):
+    if encoding.startswith("data:image/"):
+        encoding = encoding.split(";")[1].split(",")[1]
+    try:
+        image = Image.open(io.BytesIO(base64.b64decode(encoding)))
+        return image
+    except Exception as e:
+        logger.warning(f"API cannot decode image: {e}")
+        raise HTTPException(
+            status_code=500, detail="Invalid encoded image"
+        ) from e
+
+
+def to_base64_nparray(encoding: str) -> np.ndarray:
+    return np.array(decode_base64_to_image(encoding)).astype("uint8")
+
+
+class ImageMatchingAPI(torch.nn.Module):
+    default_conf = {
+        "ransac": {
+            "enable": True,
+            "estimator": "poselib",
+            "geometry": "homography",
+            "method": "RANSAC",
+            "reproj_threshold": 3,
+            "confidence": 0.9999,
+            "max_iter": 10000,
+        },
+    }
+
+    def __init__(
+        self,
+        conf: dict = {},
+        device: str = "cpu",
+        detect_threshold: float = 0.015,
+        max_keypoints: int = 1024,
+        match_threshold: float = 0.2,
+    ) -> None:
+        """
+        Initializes an instance of the ImageMatchingAPI class.
+
+        Args:
+            conf (dict): A dictionary containing the configuration parameters.
+            device (str, optional): The device to use for computation. Defaults to "cpu".
+            detect_threshold (float, optional): The threshold for detecting keypoints. Defaults to 0.015.
+            max_keypoints (int, optional): The maximum number of keypoints to extract. Defaults to 1024.
+            match_threshold (float, optional): The threshold for matching keypoints. Defaults to 0.2.
+
+        Returns:
+            None
+        """
+        super().__init__()
+        self.device = device
+        self.conf = {**self.default_conf, **conf}
+        self._updata_config(detect_threshold, max_keypoints, match_threshold)
+        self._init_models()
+        if device == "cuda":
+            memory_allocated = torch.cuda.memory_allocated(device)
+            memory_reserved = torch.cuda.memory_reserved(device)
+            logger.info(
+                f"GPU memory allocated: {memory_allocated / 1024**2:.3f} MB"
+            )
+            logger.info(
+                f"GPU memory reserved: {memory_reserved / 1024**2:.3f} MB"
+            )
+        self.pred = None
+
+    def parse_match_config(self, conf):
+        if conf["dense"]:
+            return {
+                **conf,
+                "matcher": match_dense.confs.get(
+                    conf["matcher"]["model"]["name"]
+                ),
+                "dense": True,
+            }
+        else:
+            return {
+                **conf,
+                "feature": extract_features.confs.get(
+                    conf["feature"]["model"]["name"]
+                ),
+                "matcher": match_features.confs.get(
+                    conf["matcher"]["model"]["name"]
+                ),
+                "dense": False,
+            }
+
+    def _updata_config(
+        self,
+        detect_threshold: float = 0.015,
+        max_keypoints: int = 1024,
+        match_threshold: float = 0.2,
+    ):
+        self.dense = self.conf["dense"]
+        if self.conf["dense"]:
+            try:
+                self.conf["matcher"]["model"][
+                    "match_threshold"
+                ] = match_threshold
+            except TypeError as e:
+                logger.error(e)
+        else:
+            self.conf["feature"]["model"]["max_keypoints"] = max_keypoints
+            self.conf["feature"]["model"][
+                "keypoint_threshold"
+            ] = detect_threshold
+            self.extract_conf = self.conf["feature"]
+
+        self.match_conf = self.conf["matcher"]
+
+    def _init_models(self):
+        # initialize matcher
+        self.matcher = get_model(self.match_conf)
+        # initialize extractor
+        if self.dense:
+            self.extractor = None
+        else:
+            self.extractor = get_feature_model(self.conf["feature"])
+
+    def _forward(self, img0, img1):
+        if self.dense:
+            pred = match_dense.match_images(
+                self.matcher,
+                img0,
+                img1,
+                self.match_conf["preprocessing"],
+                device=self.device,
+            )
+            last_fixed = "{}".format(  # noqa: F841
+                self.match_conf["model"]["name"]
+            )
+        else:
+            pred0 = extract_features.extract(
+                self.extractor, img0, self.extract_conf["preprocessing"]
+            )
+            pred1 = extract_features.extract(
+                self.extractor, img1, self.extract_conf["preprocessing"]
+            )
+            pred = match_features.match_images(self.matcher, pred0, pred1)
+        return pred
+
+    @torch.inference_mode()
+    def extract(self, img0: np.ndarray, **kwargs) -> Dict[str, np.ndarray]:
+        """Extract features from a single image.
+
+        Args:
+            img0 (np.ndarray): image
+
+        Returns:
+            Dict[str, np.ndarray]: feature dict
+        """
+
+        # setting prams
+        self.extractor.conf["max_keypoints"] = kwargs.get("max_keypoints", 512)
+        self.extractor.conf["keypoint_threshold"] = kwargs.get(
+            "keypoint_threshold", 0.0
+        )
+
+        pred = extract_features.extract(
+            self.extractor, img0, self.extract_conf["preprocessing"]
+        )
+        pred = {
+            k: v.cpu().detach()[0].numpy() if isinstance(v, torch.Tensor) else v
+            for k, v in pred.items()
+        }
+        # back to origin scale
+        s0 = pred["original_size"] / pred["size"]
+        pred["keypoints_orig"] = (
+            match_features.scale_keypoints(pred["keypoints"] + 0.5, s0) - 0.5
+        )
+        # TODO: rotate back
+
+        binarize = kwargs.get("binarize", False)
+        if binarize:
+            assert "descriptors" in pred
+            pred["descriptors"] = (pred["descriptors"] > 0).astype(np.uint8)
+            pred["descriptors"] = pred["descriptors"].T  # N x DIM
+        return pred
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        img0: np.ndarray,
+        img1: np.ndarray,
+    ) -> Dict[str, np.ndarray]:
+        """
+        Forward pass of the image matching API.
+
+        Args:
+            img0: A 3D NumPy array of shape (H, W, C) representing the first image.
+                  Values are in the range [0, 1] and are in RGB mode.
+            img1: A 3D NumPy array of shape (H, W, C) representing the second image.
+                  Values are in the range [0, 1] and are in RGB mode.
+
+        Returns:
+            A dictionary containing the following keys:
+            - image0_orig: The original image 0.
+            - image1_orig: The original image 1.
+            - keypoints0_orig: The keypoints detected in image 0.
+            - keypoints1_orig: The keypoints detected in image 1.
+            - mkeypoints0_orig: The raw matches between image 0 and image 1.
+            - mkeypoints1_orig: The raw matches between image 1 and image 0.
+            - mmkeypoints0_orig: The RANSAC inliers in image 0.
+            - mmkeypoints1_orig: The RANSAC inliers in image 1.
+            - mconf: The confidence scores for the raw matches.
+            - mmconf: The confidence scores for the RANSAC inliers.
+        """
+        # Take as input a pair of images (not a batch)
+        assert isinstance(img0, np.ndarray)
+        assert isinstance(img1, np.ndarray)
+        self.pred = self._forward(img0, img1)
+        if self.conf["ransac"]["enable"]:
+            self.pred = self._geometry_check(self.pred)
+        return self.pred
+
+    def _geometry_check(
+        self,
+        pred: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """
+        Filter matches using RANSAC. If keypoints are available, filter by keypoints.
+        If lines are available, filter by lines. If both keypoints and lines are
+        available, filter by keypoints.
+
+        Args:
+            pred (Dict[str, Any]): dict of matches, including original keypoints.
+                                  See :func:`filter_matches` for the expected keys.
+
+        Returns:
+            Dict[str, Any]: filtered matches
+        """
+        pred = filter_matches(
+            pred,
+            ransac_method=self.conf["ransac"]["method"],
+            ransac_reproj_threshold=self.conf["ransac"]["reproj_threshold"],
+            ransac_confidence=self.conf["ransac"]["confidence"],
+            ransac_max_iter=self.conf["ransac"]["max_iter"],
+        )
+        return pred
+
+    def visualize(
+        self,
+        log_path: Optional[Path] = None,
+    ) -> None:
+        """
+        Visualize the matches.
+
+        Args:
+            log_path (Path, optional): The directory to save the images. Defaults to None.
+
+        Returns:
+            None
+        """
+        if self.conf["dense"]:
+            postfix = str(self.conf["matcher"]["model"]["name"])
+        else:
+            postfix = "{}_{}".format(
+                str(self.conf["feature"]["model"]["name"]),
+                str(self.conf["matcher"]["model"]["name"]),
+            )
+        titles = [
+            "Image 0 - Keypoints",
+            "Image 1 - Keypoints",
+        ]
+        pred: Dict[str, Any] = self.pred
+        image0: np.ndarray = pred["image0_orig"]
+        image1: np.ndarray = pred["image1_orig"]
+        output_keypoints: np.ndarray = plot_images(
+            [image0, image1], titles=titles, dpi=300
+        )
+        if (
+            "keypoints0_orig" in pred.keys()
+            and "keypoints1_orig" in pred.keys()
+        ):
+            plot_keypoints([pred["keypoints0_orig"], pred["keypoints1_orig"]])
+            text: str = (
+                f"# keypoints0: {len(pred['keypoints0_orig'])} \n"
+                + f"# keypoints1: {len(pred['keypoints1_orig'])}"
+            )
+            add_text(0, text, fs=15)
+        output_keypoints = fig2im(output_keypoints)
+        # plot images with raw matches
+        titles = [
+            "Image 0 - Raw matched keypoints",
+            "Image 1 - Raw matched keypoints",
+        ]
+        output_matches_raw, num_matches_raw = display_matches(
+            pred, titles=titles, tag="KPTS_RAW"
+        )
+        # plot images with ransac matches
+        titles = [
+            "Image 0 - Ransac matched keypoints",
+            "Image 1 - Ransac matched keypoints",
+        ]
+        output_matches_ransac, num_matches_ransac = display_matches(
+            pred, titles=titles, tag="KPTS_RANSAC"
+        )
+        if log_path is not None:
+            img_keypoints_path: Path = log_path / f"img_keypoints_{postfix}.png"
+            img_matches_raw_path: Path = (
+                log_path / f"img_matches_raw_{postfix}.png"
+            )
+            img_matches_ransac_path: Path = (
+                log_path / f"img_matches_ransac_{postfix}.png"
+            )
+            cv2.imwrite(
+                str(img_keypoints_path),
+                output_keypoints[:, :, ::-1].copy(),  # RGB -> BGR
+            )
+            cv2.imwrite(
+                str(img_matches_raw_path),
+                output_matches_raw[:, :, ::-1].copy(),  # RGB -> BGR
+            )
+            cv2.imwrite(
+                str(img_matches_ransac_path),
+                output_matches_ransac[:, :, ::-1].copy(),  # RGB -> BGR
+            )
+            plt.close("all")
+
+
+class ImageMatchingService:
+    def __init__(self, conf: dict, device: str):
+        self.conf = conf
+        self.api = ImageMatchingAPI(conf=conf, device=device)
+        self.app = FastAPI()
+        self.register_routes()
+
+    def register_routes(self):
+
+        @self.app.get("/version")
+        async def version():
+            return {"version": get_version()}
+
+        @self.app.post("/v1/match")
+        async def match(
+            image0: UploadFile = File(...), image1: UploadFile = File(...)
+        ):
+            """
+            Handle the image matching request and return the processed result.
+
+            Args:
+                image0 (UploadFile): The first image file for matching.
+                image1 (UploadFile): The second image file for matching.
+
+            Returns:
+                JSONResponse: A JSON response containing the filtered match results
+                              or an error message in case of failure.
+            """
+            try:
+                # Load the images from the uploaded files
+                image0_array = self.load_image(image0)
+                image1_array = self.load_image(image1)
+
+                # Perform image matching using the API
+                output = self.api(image0_array, image1_array)
+
+                # Keys to skip in the output
+                skip_keys = ["image0_orig", "image1_orig"]
+
+                # Postprocess the output to filter unwanted data
+                pred = self.postprocess(output, skip_keys)
+
+                # Return the filtered prediction as a JSON response
+                return JSONResponse(content=pred)
+            except Exception as e:
+                # Return an error message with status code 500 in case of exception
+                return JSONResponse(content={"error": str(e)}, status_code=500)
+
+        @self.app.post("/v1/extract")
+        async def extract(input_info: ImagesInput):
+            """
+            Extract keypoints and descriptors from images.
+
+            Args:
+                input_info: An object containing the image data and options.
+
+            Returns:
+                A list of dictionaries containing the keypoints and descriptors.
+            """
+            try:
+                preds = []
+                for i, input_image in enumerate(input_info.data):
+                    # Load the image from the input data
+                    image_array = to_base64_nparray(input_image)
+                    # Extract keypoints and descriptors
+                    output = self.api.extract(
+                        image_array,
+                        max_keypoints=input_info.max_keypoints[i],
+                        binarize=input_info.binarize,
+                    )
+                    # Do not return the original image and image_orig
+                    # skip_keys = ["image", "image_orig"]
+                    skip_keys = []
+
+                    # Postprocess the output
+                    pred = self.postprocess(output, skip_keys)
+                    preds.append(pred)
+                # Return the list of extracted features
+                return JSONResponse(content=preds)
+            except Exception as e:
+                # Return an error message if an exception occurs
+                return JSONResponse(content={"error": str(e)}, status_code=500)
+
+    def load_image(self, file_path: Union[str, UploadFile]) -> np.ndarray:
+        """
+        Reads an image from a file path or an UploadFile object.
+
+        Args:
+            file_path: A file path or an UploadFile object.
+
+        Returns:
+            A numpy array representing the image.
+        """
+        if isinstance(file_path, str):
+            file_path = Path(file_path).resolve(strict=False)
+        else:
+            file_path = file_path.file
+        with Image.open(file_path) as img:
+            image_array = np.array(img)
+        return image_array
+
+    def postprocess(
+        self, output: dict, skip_keys: list, binarize: bool = True
+    ) -> dict:
+        pred = {}
+        for key, value in output.items():
+            if key in skip_keys:
+                continue
+            if isinstance(value, np.ndarray):
+                pred[key] = value.tolist()
+        return pred
+
+    def run(self, host: str = "0.0.0.0", port: int = 8001):
+        uvicorn.run(self.app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    conf = {
+        "feature": {
+            "output": "feats-superpoint-n4096-rmax1600",
+            "model": {
+                "name": "superpoint",
+                "nms_radius": 3,
+                "max_keypoints": 4096,
+                "keypoint_threshold": 0.005,
+            },
+            "preprocessing": {
+                "grayscale": True,
+                "force_resize": True,
+                "resize_max": 1600,
+                "width": 640,
+                "height": 480,
+                "dfactor": 8,
+            },
+        },
+        "matcher": {
+            "output": "matches-NN-mutual",
+            "model": {
+                "name": "nearest_neighbor",
+                "do_mutual_check": True,
+                "match_threshold": 0.2,
+            },
+        },
+        "dense": False,
+    }
+
+    service = ImageMatchingService(conf=conf, device=DEVICE)
+    service.run()

api/test/CMakeLists.txt

@@ -0,0 +1,16 @@
+cmake_minimum_required(VERSION 3.10)
+project(imatchui)
+
+set(OpenCV_DIR /usr/include/opencv4)
+find_package(OpenCV REQUIRED)
+
+find_package(Boost REQUIRED COMPONENTS system)
+if(Boost_FOUND)
+    include_directories(${Boost_INCLUDE_DIRS})
+endif()
+
+add_executable(client client.cpp)
+
+target_include_directories(client PRIVATE ${Boost_LIBRARIES} ${OpenCV_INCLUDE_DIRS})
+
+target_link_libraries(client PRIVATE curl jsoncpp b64 ${OpenCV_LIBS})

api/test/build_and_run.sh

@@ -0,0 +1,16 @@
+# g++ main.cpp -I/usr/include/opencv4 -lcurl -ljsoncpp -lb64 -lopencv_core -lopencv_imgcodecs -o main
+# sudo apt-get update
+# sudo apt-get install libboost-all-dev -y
+# sudo apt-get install libcurl4-openssl-dev libjsoncpp-dev libb64-dev libopencv-dev -y
+
+cd build
+cmake ..
+make -j12
+
+echo " ======== RUN DEMO ========"
+
+./client
+
+echo " ======== END DEMO ========"
+
+cd ..

api/test/client.cpp

@@ -0,0 +1,84 @@
+#include <curl/curl.h>
+#include <opencv2/opencv.hpp>
+#include "helper.h"
+
+int main() {
+    std::string img_path = "../../../datasets/sacre_coeur/mapping_rot/02928139_3448003521_rot45.jpg";
+    cv::Mat original_img = cv::imread(img_path, cv::IMREAD_GRAYSCALE);
+
+    if (original_img.empty()) {
+        throw std::runtime_error("Failed to decode image");
+    }
+
+    // Convert the image to Base64
+    std::string base64_img = image_to_base64(original_img);
+
+    // Convert the Base64 back to an image
+    cv::Mat decoded_img = base64_to_image(base64_img);
+    cv::imwrite("decoded_image.jpg", decoded_img);
+    cv::imwrite("original_img.jpg", original_img);
+
+    // The images should be identical
+    if (cv::countNonZero(original_img != decoded_img) != 0) {
+        std::cerr << "The images are not identical" << std::endl;
+        return -1;
+    } else {
+        std::cout << "The images are identical!" << std::endl;
+    }
+
+    // construct params
+    APIParams params{
+        .data = {base64_img},
+        .max_keypoints = {100, 100},
+        .timestamps = {"0", "1"},
+        .grayscale = {0},
+        .image_hw = {{480, 640}, {240, 320}},
+        .feature_type = 0,
+        .rotates = {0.0f, 0.0f},
+        .scales = {1.0f, 1.0f},
+        .reference_points = {
+            {1.23e+2f, 1.2e+1f},
+            {5.0e-1f, 3.0e-1f},
+            {2.3e+2f, 2.2e+1f},
+            {6.0e-1f, 4.0e-1f}
+        },
+        .binarize = {1}
+    };
+
+    KeyPointResults kpts_results;
+
+    // Convert the parameters to JSON
+    Json::Value jsonData = paramsToJson(params);
+    std::string url = "http://127.0.0.1:8001/v1/extract";
+    Json::StreamWriterBuilder writer;
+    std::string output = Json::writeString(writer, jsonData);
+
+    CURL* curl;
+    CURLcode res;
+    std::string readBuffer;
+
+    curl_global_init(CURL_GLOBAL_DEFAULT);
+    curl = curl_easy_init();
+    if (curl) {
+        struct curl_slist* hs = NULL;
+        hs = curl_slist_append(hs, "Content-Type: application/json");
+        curl_easy_setopt(curl, CURLOPT_HTTPHEADER, hs);
+        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
+        curl_easy_setopt(curl, CURLOPT_POSTFIELDS, output.c_str());
+        curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, WriteCallback);
+        curl_easy_setopt(curl, CURLOPT_WRITEDATA, &readBuffer);
+        res = curl_easy_perform(curl);
+
+        if (res != CURLE_OK)
+            fprintf(stderr, "curl_easy_perform() failed: %s\n",
+                    curl_easy_strerror(res));
+        else {
+            // std::cout << "Response from server: " << readBuffer << std::endl;
+           kpts_results =  decode_response(readBuffer);
+        }
+        curl_easy_cleanup(curl);
+    }
+    curl_global_cleanup();
+
+    return 0;
+}

api/test/helper.h

@@ -0,0 +1,410 @@
+
+#include <sstream>
+#include <fstream>
+#include <vector>
+#include <b64/encode.h>
+#include <jsoncpp/json/json.h>
+#include <opencv2/opencv.hpp>
+
+// base64 to image
+#include <boost/archive/iterators/binary_from_base64.hpp>
+#include <boost/archive/iterators/transform_width.hpp>
+#include <boost/archive/iterators/base64_from_binary.hpp>
+
+/// Parameters used in the API
+struct APIParams {
+    /// A list of images, base64 encoded
+    std::vector<std::string> data;
+
+    /// The maximum number of keypoints to detect for each image
+    std::vector<int> max_keypoints;
+
+    /// The timestamps of the images
+    std::vector<std::string> timestamps;
+
+    /// Whether to convert the images to grayscale
+    bool grayscale;
+
+    /// The height and width of each image
+    std::vector<std::vector<int>> image_hw;
+
+    /// The type of feature detector to use
+    int feature_type;
+
+    /// The rotations of the images
+    std::vector<double> rotates;
+
+    /// The scales of the images
+    std::vector<double> scales;
+
+    /// The reference points of the images
+    std::vector<std::vector<float>> reference_points;
+
+    /// Whether to binarize the descriptors
+    bool binarize;
+};
+
+/**
+ * @brief Contains the results of a keypoint detector.
+ *
+ * @details Stores the keypoints and descriptors for each image.
+ */
+class KeyPointResults {
+public:
+    KeyPointResults() {}
+
+    /**
+     * @brief Constructor.
+     *
+     * @param kp The keypoints for each image.
+     */
+    KeyPointResults(const std::vector<std::vector<cv::KeyPoint>>& kp, 
+                    const std::vector<cv::Mat>& desc)
+        : keypoints(kp), descriptors(desc) {}
+
+    /**
+     * @brief Append keypoints to the result.
+     *
+     * @param kpts The keypoints to append.
+     */
+    inline void append_keypoints(std::vector<cv::KeyPoint>& kpts) {
+        keypoints.emplace_back(kpts);
+    }
+
+    /**
+     * @brief Append descriptors to the result.
+     *
+     * @param desc The descriptors to append.
+     */
+    inline void append_descriptors(cv::Mat& desc) {
+        descriptors.emplace_back(desc);
+    }
+
+    /**
+     * @brief Get the keypoints.
+     *
+     * @return The keypoints.
+     */
+    inline std::vector<std::vector<cv::KeyPoint>> get_keypoints() {
+        return keypoints;
+    }
+
+    /**
+     * @brief Get the descriptors.
+     *
+     * @return The descriptors.
+     */
+    inline std::vector<cv::Mat> get_descriptors() {
+        return descriptors;
+    }
+
+private:
+    std::vector<std::vector<cv::KeyPoint>> keypoints;
+    std::vector<cv::Mat> descriptors;
+    std::vector<std::vector<float>> scores;
+};
+
+
+/**
+ * @brief Decodes a base64 encoded string.
+ *
+ * @param base64 The base64 encoded string to decode.
+ * @return The decoded string.
+ */
+std::string base64_decode(const std::string& base64) {
+    using namespace boost::archive::iterators;
+    using It = transform_width<binary_from_base64<std::string::const_iterator>, 8, 6>;
+
+    // Find the position of the last non-whitespace character
+    auto end = base64.find_last_not_of(" \t\n\r");
+    if (end != std::string::npos) {
+        // Move one past the last non-whitespace character
+        end += 1;
+    }
+
+    // Decode the base64 string and return the result
+    return std::string(It(base64.begin()), It(base64.begin() + end));
+}
+
+
+
+/**
+ * @brief Decodes a base64 string into an OpenCV image
+ *
+ * @param base64 The base64 encoded string
+ * @return The decoded OpenCV image
+ */
+cv::Mat base64_to_image(const std::string& base64) {
+    // Decode the base64 string
+    std::string decodedStr = base64_decode(base64);
+
+    // Decode the image
+    std::vector<uchar> data(decodedStr.begin(), decodedStr.end());
+    cv::Mat img = cv::imdecode(data, cv::IMREAD_GRAYSCALE);
+
+    // Check for errors
+    if (img.empty()) {
+        throw std::runtime_error("Failed to decode image");
+    }
+
+    return img;
+}
+
+
+/**
+ * @brief Encodes an OpenCV image into a base64 string
+ *
+ * This function takes an OpenCV image and encodes it into a base64 string.
+ * The image is first encoded as a PNG image, and then the resulting
+ * bytes are encoded as a base64 string.
+ *
+ * @param img The OpenCV image
+ * @return The base64 encoded string
+ *
+ * @throws std::runtime_error if the image is empty or encoding fails
+ */
+std::string image_to_base64(cv::Mat &img) {
+    if (img.empty()) {
+        throw std::runtime_error("Failed to read image");
+    }
+
+    // Encode the image as a PNG
+    std::vector<uchar> buf;
+    if (!cv::imencode(".png", img, buf)) {
+        throw std::runtime_error("Failed to encode image");
+    }
+
+    // Encode the bytes as a base64 string
+    using namespace boost::archive::iterators;
+    using It = base64_from_binary<transform_width<std::vector<uchar>::const_iterator, 6, 8>>;
+    std::string base64(It(buf.begin()), It(buf.end()));
+
+    // Pad the string with '=' characters to a multiple of 4 bytes
+    base64.append((3 - buf.size() % 3) % 3, '=');
+
+    return base64;
+}
+
+
+/**
+ * @brief Callback function for libcurl to write data to a string
+ *
+ * This function is used as a callback for libcurl to write data to a string.
+ * It takes the contents, size, and nmemb as parameters, and writes the data to
+ * the string.
+ *
+ * @param contents The data to write
+ * @param size The size of the data
+ * @param nmemb The number of members in the data
+ * @param s The string to write the data to
+ * @return The number of bytes written
+ */
+size_t WriteCallback(void* contents, size_t size, size_t nmemb, std::string* s) {
+    size_t newLength = size * nmemb;
+    try {
+        // Resize the string to fit the new data
+        s->resize(s->size() + newLength);
+    } catch (std::bad_alloc& e) {
+        // If there's an error allocating memory, return 0
+        return 0;
+    }
+
+    // Copy the data to the string
+    std::copy(static_cast<const char*>(contents),
+              static_cast<const char*>(contents) + newLength,
+              s->begin() + s->size() - newLength);
+    return newLength;
+}
+
+// Helper functions
+
+/**
+ * @brief Helper function to convert a type to a Json::Value
+ *
+ * This function takes a value of type T and converts it to a Json::Value.
+ * It is used to simplify the process of converting a type to a Json::Value.
+ *
+ * @param val The value to convert
+ * @return The converted Json::Value
+ */
+template <typename T>
+Json::Value toJson(const T& val) {
+    return Json::Value(val);
+}
+
+/**
+ * @brief Converts a vector to a Json::Value
+ *
+ * This function takes a vector of type T and converts it to a Json::Value.
+ * Each element in the vector is appended to the Json::Value array.
+ *
+ * @param vec The vector to convert to Json::Value
+ * @return The Json::Value representing the vector
+ */
+template <typename T>
+Json::Value vectorToJson(const std::vector<T>& vec) {
+    Json::Value json(Json::arrayValue);
+    for (const auto& item : vec) {
+        json.append(item);
+    }
+    return json;
+}
+
+/**
+ * @brief Converts a nested vector to a Json::Value
+ *
+ * This function takes a nested vector of type T and converts it to a Json::Value.
+ * Each sub-vector is converted to a Json::Value array and appended to the main Json::Value array.
+ *
+ * @param vec The nested vector to convert to Json::Value
+ * @return The Json::Value representing the nested vector
+ */
+template <typename T>
+Json::Value nestedVectorToJson(const std::vector<std::vector<T>>& vec) {
+    Json::Value json(Json::arrayValue);
+    for (const auto& subVec : vec) {
+        json.append(vectorToJson(subVec));
+    }
+    return json;
+}
+
+
+
+/**
+ * @brief Converts the APIParams struct to a Json::Value
+ *
+ * This function takes an APIParams struct and converts it to a Json::Value.
+ * The Json::Value is a JSON object with the following fields:
+ * - data: a JSON array of base64 encoded images
+ * - max_keypoints: a JSON array of integers, max number of keypoints for each image
+ * - timestamps: a JSON array of timestamps, one for each image
+ * - grayscale: a JSON boolean, whether to convert images to grayscale
+ * - image_hw: a nested JSON array, each sub-array contains the height and width of an image
+ * - feature_type: a JSON integer, the type of feature detector to use
+ * - rotates: a JSON array of doubles, the rotation of each image
+ * - scales: a JSON array of doubles, the scale of each image
+ * - reference_points: a nested JSON array, each sub-array contains the reference points of an image
+ * - binarize: a JSON boolean, whether to binarize the descriptors
+ *
+ * @param params The APIParams struct to convert
+ * @return The Json::Value representing the APIParams struct
+ */
+Json::Value paramsToJson(const APIParams& params) {
+    Json::Value json;
+    json["data"] = vectorToJson(params.data);
+    json["max_keypoints"] = vectorToJson(params.max_keypoints);
+    json["timestamps"] = vectorToJson(params.timestamps);
+    json["grayscale"] = toJson(params.grayscale);
+    json["image_hw"] = nestedVectorToJson(params.image_hw);
+    json["feature_type"] = toJson(params.feature_type);
+    json["rotates"] = vectorToJson(params.rotates);
+    json["scales"] = vectorToJson(params.scales);
+    json["reference_points"] = nestedVectorToJson(params.reference_points);
+    json["binarize"] = toJson(params.binarize);
+    return json;
+}
+
+template<typename T>
+cv::Mat jsonToMat(Json::Value json) {
+    int rows = json.size();
+    int cols = json[0].size();
+
+    // Create a single array to hold all the data.
+    std::vector<T> data;
+    data.reserve(rows * cols);
+
+    for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < cols; j++) {
+            data.push_back(static_cast<T>(json[i][j].asInt()));
+        }
+    }
+
+    // Create a cv::Mat object that points to the data.
+    cv::Mat mat(rows, cols, CV_8UC1, data.data());  // Change the type if necessary.
+    // cv::Mat mat(cols, rows,CV_8UC1, data.data());  // Change the type if necessary.
+
+    return mat;
+}
+
+
+
+/**
+ * @brief Decodes the response of the server and prints the keypoints
+ *
+ * This function takes the response of the server, a JSON string, and decodes
+ * it. It then prints the keypoints and draws them on the original image.
+ *
+ * @param response The response of the server
+ * @return The keypoints and descriptors
+ */
+KeyPointResults decode_response(const std::string& response, bool viz=true) {
+    Json::CharReaderBuilder builder;
+    Json::CharReader* reader = builder.newCharReader();
+
+    Json::Value jsonData;
+    std::string errors;
+
+    // Parse the JSON response
+    bool parsingSuccessful = reader->parse(response.c_str(),
+        response.c_str() + response.size(), &jsonData, &errors);
+    delete reader;
+
+    if (!parsingSuccessful) {
+        // Handle error
+        std::cout << "Failed to parse the JSON, errors:" << std::endl;
+        std::cout << errors << std::endl;
+        return KeyPointResults();
+    }
+
+    KeyPointResults kpts_results;
+
+    // Iterate over the images
+    for (const auto& jsonItem : jsonData) {
+        auto jkeypoints = jsonItem["keypoints"];
+        auto jkeypoints_orig = jsonItem["keypoints_orig"];
+        auto jdescriptors = jsonItem["descriptors"];
+        auto jscores = jsonItem["scores"];
+        auto jimageSize = jsonItem["image_size"];
+        auto joriginalSize = jsonItem["original_size"];
+        auto jsize = jsonItem["size"];
+
+        std::vector<cv::KeyPoint> vkeypoints;
+        std::vector<float> vscores;
+
+        // Iterate over the keypoints
+        int counter = 0;
+        for (const auto& keypoint : jkeypoints_orig) {
+            if (counter < 10) {
+                // Print the first 10 keypoints
+                std::cout << keypoint[0].asFloat() << ", "
+                    << keypoint[1].asFloat() << std::endl;
+            }
+            counter++;
+            // Convert the Json::Value to a cv::KeyPoint
+            vkeypoints.emplace_back(cv::KeyPoint(keypoint[0].asFloat(),
+                keypoint[1].asFloat(), 0.0));
+        }
+
+        if (viz && jsonItem.isMember("image_orig")) {
+
+            auto jimg_orig = jsonItem["image_orig"];
+            cv::Mat img = jsonToMat<uchar>(jimg_orig);
+            cv::imwrite("viz_image_orig.jpg", img);
+
+            // Draw keypoints on the image
+            cv::Mat imgWithKeypoints;
+            cv::drawKeypoints(img, vkeypoints, 
+                imgWithKeypoints, cv::Scalar(0, 0, 255));
+
+            // Write the image with keypoints
+            std::string filename = "viz_image_orig_keypoints.jpg";
+            cv::imwrite(filename, imgWithKeypoints);
+        }
+
+        // Iterate over the descriptors
+        cv::Mat descriptors = jsonToMat<uchar>(jdescriptors); 
+        kpts_results.append_keypoints(vkeypoints);
+        kpts_results.append_descriptors(descriptors);
+    }
+    return kpts_results;
+}

api/types.py

@@ -0,0 +1,16 @@
+from typing import List
+
+from pydantic import BaseModel
+
+
+class ImagesInput(BaseModel):
+    data: List[str] = []
+    max_keypoints: List[int] = []
+    timestamps: List[str] = []
+    grayscale: bool = False
+    image_hw: List[List[int]] = [[], []]
+    feature_type: int = 0
+    rotates: List[float] = []
+    scales: List[float] = []
+    reference_points: List[List[float]] = []
+    binarize: bool = False

app.py

@@ -0,0 +1,28 @@
+import argparse
+from pathlib import Path
+from ui.app_class import ImageMatchingApp
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--server_name",
+        type=str,
+        default="0.0.0.0",
+        help="server name",
+    )
+    parser.add_argument(
+        "--server_port",
+        type=int,
+        default=7860,
+        help="server port",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=Path(__file__).parent / "ui/config.yaml",
+        help="config file",
+    )
+    args = parser.parse_args()
+    ImageMatchingApp(
+        args.server_name, args.server_port, config=args.config
+    ).run()

build_docker.sh

@@ -0,0 +1,3 @@
+docker build -t image-matching-webui:latest . --no-cache
+docker tag image-matching-webui:latest vincentqin/image-matching-webui:latest
+docker push vincentqin/image-matching-webui:latest

docker/Dockerfile

@@ -0,0 +1,27 @@
+# Use an official conda-based Python image as a parent image
+FROM pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime
+LABEL maintainer vincentqyw
+ARG PYTHON_VERSION=3.10.10
+
+# Set the working directory to /code
+WORKDIR /code
+
+# Install Git and Git LFS
+RUN apt-get update && apt-get install -y git-lfs
+RUN git lfs install
+
+# Clone the Git repository
+RUN git clone https://huggingface.co/spaces/Realcat/image-matching-webui /code
+
+RUN conda create -n imw python=${PYTHON_VERSION}
+RUN echo "source activate imw" > ~/.bashrc
+ENV PATH /opt/conda/envs/imw/bin:$PATH
+
+# Make RUN commands use the new environment
+SHELL ["conda", "run", "-n", "imw", "/bin/bash", "-c"]
+RUN pip install --upgrade pip
+RUN pip install -r requirements.txt
+RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
+
+# Export port
+EXPOSE 7860

docker/build_docker.bat

@@ -0,0 +1,3 @@
+docker build -t image-matching-webui:latest . --no-cache
+# docker tag image-matching-webui:latest vincentqin/image-matching-webui:latest
+# docker push vincentqin/image-matching-webui:latest

docker/run_docker.bat

@@ -0,0 +1 @@
+docker run -it -p 7860:7860 vincentqin/image-matching-webui:latest python app.py --server_name "0.0.0.0" --server_port=7860

docker/run_docker.sh

@@ -0,0 +1 @@
+docker run -it -p 7860:7860 vincentqin/image-matching-webui:latest python app.py --server_name "0.0.0.0" --server_port=7860

format.sh

@@ -0,0 +1,3 @@
+python -m flake8 ui/*.py api/*.py hloc/*.py hloc/matchers/*.py hloc/extractors/*.py 
+python -m isort ui/*.py api/*.py hloc/*.py hloc/matchers/*.py hloc/extractors/*.py
+python -m black ui/*.py api/*.py hloc/*.py hloc/matchers/*.py hloc/extractors/*.py

hloc/__init__.py

@@ -0,0 +1,63 @@
+import logging
+import sys
+
+import torch
+from packaging import version
+
+__version__ = "1.5"
+
+LOG_PATH = "log.txt"
+
+
+def read_logs():
+    sys.stdout.flush()
+    with open(LOG_PATH, "r") as f:
+        return f.read()
+
+
+def flush_logs():
+    sys.stdout.flush()
+    logs = open(LOG_PATH, "w")
+    logs.close()
+
+
+formatter = logging.Formatter(
+    fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s",
+    datefmt="%Y/%m/%d %H:%M:%S",
+)
+
+logs_file = open(LOG_PATH, "w")
+logs_file.close()
+
+file_handler = logging.FileHandler(filename=LOG_PATH)
+file_handler.setFormatter(formatter)
+file_handler.setLevel(logging.INFO)
+stdout_handler = logging.StreamHandler()
+stdout_handler.setFormatter(formatter)
+stdout_handler.setLevel(logging.INFO)
+logger = logging.getLogger("hloc")
+logger.setLevel(logging.INFO)
+logger.addHandler(file_handler)
+logger.addHandler(stdout_handler)
+logger.propagate = False
+
+try:
+    import pycolmap
+except ImportError:
+    logger.warning("pycolmap is not installed, some features may not work.")
+else:
+    min_version = version.parse("0.6.0")
+    found_version = pycolmap.__version__
+    if found_version != "dev":
+        version = version.parse(found_version)
+        if version < min_version:
+            s = f"pycolmap>={min_version}"
+            logger.warning(
+                "hloc requires %s but found pycolmap==%s, "
+                'please upgrade with `pip install --upgrade "%s"`',
+                s,
+                found_version,
+                s,
+            )
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

hloc/colmap_from_nvm.py

@@ -0,0 +1,220 @@
+import argparse
+import sqlite3
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+from tqdm import tqdm
+
+from . import logger
+from .utils.read_write_model import (
+    CAMERA_MODEL_NAMES,
+    Camera,
+    Image,
+    Point3D,
+    write_model,
+)
+
+
+def recover_database_images_and_ids(database_path):
+    images = {}
+    cameras = {}
+    db = sqlite3.connect(str(database_path))
+    ret = db.execute("SELECT name, image_id, camera_id FROM images;")
+    for name, image_id, camera_id in ret:
+        images[name] = image_id
+        cameras[name] = camera_id
+    db.close()
+    logger.info(
+        f"Found {len(images)} images and {len(cameras)} cameras in database."
+    )
+    return images, cameras
+
+
+def quaternion_to_rotation_matrix(qvec):
+    qvec = qvec / np.linalg.norm(qvec)
+    w, x, y, z = qvec
+    R = np.array(
+        [
+            [
+                1 - 2 * y * y - 2 * z * z,
+                2 * x * y - 2 * z * w,
+                2 * x * z + 2 * y * w,
+            ],
+            [
+                2 * x * y + 2 * z * w,
+                1 - 2 * x * x - 2 * z * z,
+                2 * y * z - 2 * x * w,
+            ],
+            [
+                2 * x * z - 2 * y * w,
+                2 * y * z + 2 * x * w,
+                1 - 2 * x * x - 2 * y * y,
+            ],
+        ]
+    )
+    return R
+
+
+def camera_center_to_translation(c, qvec):
+    R = quaternion_to_rotation_matrix(qvec)
+    return (-1) * np.matmul(R, c)
+
+
+def read_nvm_model(
+    nvm_path, intrinsics_path, image_ids, camera_ids, skip_points=False
+):
+    with open(intrinsics_path, "r") as f:
+        raw_intrinsics = f.readlines()
+
+    logger.info(f"Reading {len(raw_intrinsics)} cameras...")
+    cameras = {}
+    for intrinsics in raw_intrinsics:
+        intrinsics = intrinsics.strip("\n").split(" ")
+        name, camera_model, width, height = intrinsics[:4]
+        params = [float(p) for p in intrinsics[4:]]
+        camera_model = CAMERA_MODEL_NAMES[camera_model]
+        assert len(params) == camera_model.num_params
+        camera_id = camera_ids[name]
+        camera = Camera(
+            id=camera_id,
+            model=camera_model.model_name,
+            width=int(width),
+            height=int(height),
+            params=params,
+        )
+        cameras[camera_id] = camera
+
+    nvm_f = open(nvm_path, "r")
+    line = nvm_f.readline()
+    while line == "\n" or line.startswith("NVM_V3"):
+        line = nvm_f.readline()
+    num_images = int(line)
+    assert num_images == len(cameras)
+
+    logger.info(f"Reading {num_images} images...")
+    image_idx_to_db_image_id = []
+    image_data = []
+    i = 0
+    while i < num_images:
+        line = nvm_f.readline()
+        if line == "\n":
+            continue
+        data = line.strip("\n").split(" ")
+        image_data.append(data)
+        image_idx_to_db_image_id.append(image_ids[data[0]])
+        i += 1
+
+    line = nvm_f.readline()
+    while line == "\n":
+        line = nvm_f.readline()
+    num_points = int(line)
+
+    if skip_points:
+        logger.info(f"Skipping {num_points} points.")
+        num_points = 0
+    else:
+        logger.info(f"Reading {num_points} points...")
+    points3D = {}
+    image_idx_to_keypoints = defaultdict(list)
+    i = 0
+    pbar = tqdm(total=num_points, unit="pts")
+    while i < num_points:
+        line = nvm_f.readline()
+        if line == "\n":
+            continue
+
+        data = line.strip("\n").split(" ")
+        x, y, z, r, g, b, num_observations = data[:7]
+        obs_image_ids, point2D_idxs = [], []
+        for j in range(int(num_observations)):
+            s = 7 + 4 * j
+            img_index, kp_index, kx, ky = data[s : s + 4]
+            image_idx_to_keypoints[int(img_index)].append(
+                (int(kp_index), float(kx), float(ky), i)
+            )
+            db_image_id = image_idx_to_db_image_id[int(img_index)]
+            obs_image_ids.append(db_image_id)
+            point2D_idxs.append(kp_index)
+
+        point = Point3D(
+            id=i,
+            xyz=np.array([x, y, z], float),
+            rgb=np.array([r, g, b], int),
+            error=1.0,  # fake
+            image_ids=np.array(obs_image_ids, int),
+            point2D_idxs=np.array(point2D_idxs, int),
+        )
+        points3D[i] = point
+
+        i += 1
+        pbar.update(1)
+    pbar.close()
+
+    logger.info("Parsing image data...")
+    images = {}
+    for i, data in enumerate(image_data):
+        # Skip the focal length. Skip the distortion and terminal 0.
+        name, _, qw, qx, qy, qz, cx, cy, cz, _, _ = data
+        qvec = np.array([qw, qx, qy, qz], float)
+        c = np.array([cx, cy, cz], float)
+        t = camera_center_to_translation(c, qvec)
+
+        if i in image_idx_to_keypoints:
+            # NVM only stores triangulated 2D keypoints: add dummy ones
+            keypoints = image_idx_to_keypoints[i]
+            point2D_idxs = np.array([d[0] for d in keypoints])
+            tri_xys = np.array([[x, y] for _, x, y, _ in keypoints])
+            tri_ids = np.array([i for _, _, _, i in keypoints])
+
+            num_2Dpoints = max(point2D_idxs) + 1
+            xys = np.zeros((num_2Dpoints, 2), float)
+            point3D_ids = np.full(num_2Dpoints, -1, int)
+            xys[point2D_idxs] = tri_xys
+            point3D_ids[point2D_idxs] = tri_ids
+        else:
+            xys = np.zeros((0, 2), float)
+            point3D_ids = np.full(0, -1, int)
+
+        image_id = image_ids[name]
+        image = Image(
+            id=image_id,
+            qvec=qvec,
+            tvec=t,
+            camera_id=camera_ids[name],
+            name=name,
+            xys=xys,
+            point3D_ids=point3D_ids,
+        )
+        images[image_id] = image
+
+    return cameras, images, points3D
+
+
+def main(nvm, intrinsics, database, output, skip_points=False):
+    assert nvm.exists(), nvm
+    assert intrinsics.exists(), intrinsics
+    assert database.exists(), database
+
+    image_ids, camera_ids = recover_database_images_and_ids(database)
+
+    logger.info("Reading the NVM model...")
+    model = read_nvm_model(
+        nvm, intrinsics, image_ids, camera_ids, skip_points=skip_points
+    )
+
+    logger.info("Writing the COLMAP model...")
+    output.mkdir(exist_ok=True, parents=True)
+    write_model(*model, path=str(output), ext=".bin")
+    logger.info("Done.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--nvm", required=True, type=Path)
+    parser.add_argument("--intrinsics", required=True, type=Path)
+    parser.add_argument("--database", required=True, type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument("--skip_points", action="store_true")
+    args = parser.parse_args()
+    main(**args.__dict__)

hloc/extract_features.py

@@ -0,0 +1,618 @@
+import argparse
+import collections.abc as collections
+import pprint
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Dict, List, Optional, Union
+
+import cv2
+import h5py
+import numpy as np
+import PIL.Image
+import torch
+import torchvision.transforms.functional as F
+from tqdm import tqdm
+
+from . import extractors, logger
+from .utils.base_model import dynamic_load
+from .utils.io import list_h5_names, read_image
+from .utils.parsers import parse_image_lists
+
+"""
+A set of standard configurations that can be directly selected from the command
+line using their name. Each is a dictionary with the following entries:
+    - output: the name of the feature file that will be generated.
+    - model: the model configuration, as passed to a feature extractor.
+    - preprocessing: how to preprocess the images read from disk.
+"""
+confs = {
+    "superpoint_aachen": {
+        "output": "feats-superpoint-n4096-r1024",
+        "model": {
+            "name": "superpoint",
+            "nms_radius": 3,
+            "max_keypoints": 4096,
+            "keypoint_threshold": 0.005,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1600,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    # Resize images to 1600px even if they are originally smaller.
+    # Improves the keypoint localization if the images are of good quality.
+    "superpoint_max": {
+        "output": "feats-superpoint-n4096-rmax1600",
+        "model": {
+            "name": "superpoint",
+            "nms_radius": 3,
+            "max_keypoints": 4096,
+            "keypoint_threshold": 0.005,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1600,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "superpoint_inloc": {
+        "output": "feats-superpoint-n4096-r1600",
+        "model": {
+            "name": "superpoint",
+            "nms_radius": 4,
+            "max_keypoints": 4096,
+            "keypoint_threshold": 0.005,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "r2d2": {
+        "output": "feats-r2d2-n5000-r1024",
+        "model": {
+            "name": "r2d2",
+            "max_keypoints": 5000,
+            "reliability_threshold": 0.7,
+            "repetability_threshold": 0.7,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "d2net-ss": {
+        "output": "feats-d2net-ss-n5000-r1600",
+        "model": {
+            "name": "d2net",
+            "multiscale": False,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "d2net-ms": {
+        "output": "feats-d2net-ms-n5000-r1600",
+        "model": {
+            "name": "d2net",
+            "multiscale": True,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "rord": {
+        "output": "feats-rord-ss-n5000-r1600",
+        "model": {
+            "name": "rord",
+            "multiscale": False,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "rootsift": {
+        "output": "feats-rootsift-n5000-r1600",
+        "model": {
+            "name": "dog",
+            "descriptor": "rootsift",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1600,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "sift": {
+        "output": "feats-sift-n5000-r1600",
+        "model": {
+            "name": "sift",
+            "rootsift": True,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1600,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "sosnet": {
+        "output": "feats-sosnet-n5000-r1600",
+        "model": {
+            "name": "dog",
+            "descriptor": "sosnet",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "hardnet": {
+        "output": "feats-hardnet-n5000-r1600",
+        "model": {
+            "name": "dog",
+            "descriptor": "hardnet",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "disk": {
+        "output": "feats-disk-n5000-r1600",
+        "model": {
+            "name": "disk",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "xfeat": {
+        "output": "feats-xfeat-n5000-r1600",
+        "model": {
+            "name": "xfeat",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "alike": {
+        "output": "feats-alike-n5000-r1600",
+        "model": {
+            "name": "alike",
+            "max_keypoints": 5000,
+            "use_relu": True,
+            "multiscale": False,
+            "detection_threshold": 0.5,
+            "top_k": -1,
+            "sub_pixel": False,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "lanet": {
+        "output": "feats-lanet-n5000-r1600",
+        "model": {
+            "name": "lanet",
+            "keypoint_threshold": 0.1,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+            "force_resize": True,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "darkfeat": {
+        "output": "feats-darkfeat-n5000-r1600",
+        "model": {
+            "name": "darkfeat",
+            "max_keypoints": 5000,
+            "reliability_threshold": 0.7,
+            "repetability_threshold": 0.7,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1600,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "dedode": {
+        "output": "feats-dedode-n5000-r1600",
+        "model": {
+            "name": "dedode",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1600,
+            "width": 768,
+            "height": 768,
+            "dfactor": 8,
+        },
+    },
+    "example": {
+        "output": "feats-example-n2000-r1024",
+        "model": {
+            "name": "example",
+            "keypoint_threshold": 0.1,
+            "max_keypoints": 2000,
+            "model_name": "model.pth",
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 768,
+            "height": 768,
+            "dfactor": 8,
+        },
+    },
+    "sfd2": {
+        "output": "feats-sfd2-n4096-r1600",
+        "model": {
+            "name": "sfd2",
+            "max_keypoints": 4096,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1600,
+            "width": 640,
+            "height": 480,
+            "conf_th": 0.001,
+            "multiscale": False,
+            "scales": [1.0],
+        },
+    },
+    # Global descriptors
+    "dir": {
+        "output": "global-feats-dir",
+        "model": {"name": "dir"},
+        "preprocessing": {"resize_max": 1024},
+    },
+    "netvlad": {
+        "output": "global-feats-netvlad",
+        "model": {"name": "netvlad"},
+        "preprocessing": {"resize_max": 1024},
+    },
+    "openibl": {
+        "output": "global-feats-openibl",
+        "model": {"name": "openibl"},
+        "preprocessing": {"resize_max": 1024},
+    },
+    "cosplace": {
+        "output": "global-feats-cosplace",
+        "model": {"name": "cosplace"},
+        "preprocessing": {"resize_max": 1024},
+    },
+    "eigenplaces": {
+        "output": "global-feats-eigenplaces",
+        "model": {"name": "eigenplaces"},
+        "preprocessing": {"resize_max": 1024},
+    },
+}
+
+
+def resize_image(image, size, interp):
+    if interp.startswith("cv2_"):
+        interp = getattr(cv2, "INTER_" + interp[len("cv2_") :].upper())
+        h, w = image.shape[:2]
+        if interp == cv2.INTER_AREA and (w < size[0] or h < size[1]):
+            interp = cv2.INTER_LINEAR
+        resized = cv2.resize(image, size, interpolation=interp)
+    elif interp.startswith("pil_"):
+        interp = getattr(PIL.Image, interp[len("pil_") :].upper())
+        resized = PIL.Image.fromarray(image.astype(np.uint8))
+        resized = resized.resize(size, resample=interp)
+        resized = np.asarray(resized, dtype=image.dtype)
+    else:
+        raise ValueError(f"Unknown interpolation {interp}.")
+    return resized
+
+
+class ImageDataset(torch.utils.data.Dataset):
+    default_conf = {
+        "globs": ["*.jpg", "*.png", "*.jpeg", "*.JPG", "*.PNG"],
+        "grayscale": False,
+        "resize_max": None,
+        "force_resize": False,
+        "interpolation": "cv2_area",  # pil_linear is more accurate but slower
+    }
+
+    def __init__(self, root, conf, paths=None):
+        self.conf = conf = SimpleNamespace(**{**self.default_conf, **conf})
+        self.root = root
+
+        if paths is None:
+            paths = []
+            for g in conf.globs:
+                paths += list(Path(root).glob("**/" + g))
+            if len(paths) == 0:
+                raise ValueError(f"Could not find any image in root: {root}.")
+            paths = sorted(list(set(paths)))
+            self.names = [i.relative_to(root).as_posix() for i in paths]
+            logger.info(f"Found {len(self.names)} images in root {root}.")
+        else:
+            if isinstance(paths, (Path, str)):
+                self.names = parse_image_lists(paths)
+            elif isinstance(paths, collections.Iterable):
+                self.names = [
+                    p.as_posix() if isinstance(p, Path) else p for p in paths
+                ]
+            else:
+                raise ValueError(f"Unknown format for path argument {paths}.")
+
+            for name in self.names:
+                if not (root / name).exists():
+                    raise ValueError(
+                        f"Image {name} does not exists in root: {root}."
+                    )
+
+    def __getitem__(self, idx):
+        name = self.names[idx]
+        image = read_image(self.root / name, self.conf.grayscale)
+        image = image.astype(np.float32)
+        size = image.shape[:2][::-1]
+
+        if self.conf.resize_max and (
+            self.conf.force_resize or max(size) > self.conf.resize_max
+        ):
+            scale = self.conf.resize_max / max(size)
+            size_new = tuple(int(round(x * scale)) for x in size)
+            image = resize_image(image, size_new, self.conf.interpolation)
+
+        if self.conf.grayscale:
+            image = image[None]
+        else:
+            image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+        image = image / 255.0
+
+        data = {
+            "image": image,
+            "original_size": np.array(size),
+        }
+        return data
+
+    def __len__(self):
+        return len(self.names)
+
+
+def extract(model, image_0, conf):
+    default_conf = {
+        "grayscale": True,
+        "resize_max": 1024,
+        "dfactor": 8,
+        "cache_images": False,
+        "force_resize": False,
+        "width": 320,
+        "height": 240,
+        "interpolation": "cv2_area",
+    }
+    conf = SimpleNamespace(**{**default_conf, **conf})
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def preprocess(image: np.ndarray, conf: SimpleNamespace):
+        image = image.astype(np.float32, copy=False)
+        size = image.shape[:2][::-1]
+        scale = np.array([1.0, 1.0])
+        if conf.resize_max:
+            scale = conf.resize_max / max(size)
+            if scale < 1.0:
+                size_new = tuple(int(round(x * scale)) for x in size)
+                image = resize_image(image, size_new, "cv2_area")
+                scale = np.array(size) / np.array(size_new)
+        if conf.force_resize:
+            image = resize_image(image, (conf.width, conf.height), "cv2_area")
+            size_new = (conf.width, conf.height)
+            scale = np.array(size) / np.array(size_new)
+        if conf.grayscale:
+            assert image.ndim == 2, image.shape
+            image = image[None]
+        else:
+            image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+        image = torch.from_numpy(image / 255.0).float()
+
+        # assure that the size is divisible by dfactor
+        size_new = tuple(
+            map(
+                lambda x: int(x // conf.dfactor * conf.dfactor),
+                image.shape[-2:],
+            )
+        )
+        image = F.resize(image, size=size_new, antialias=True)
+        input_ = image.to(device, non_blocking=True)[None]
+        data = {
+            "image": input_,
+            "image_orig": image_0,
+            "original_size": np.array(size),
+            "size": np.array(image.shape[1:][::-1]),
+        }
+        return data
+
+    # convert to grayscale if needed
+    if len(image_0.shape) == 3 and conf.grayscale:
+        image0 = cv2.cvtColor(image_0, cv2.COLOR_RGB2GRAY)
+    else:
+        image0 = image_0
+    # comment following lines, image is always RGB mode
+    # if not conf.grayscale and len(image_0.shape) == 3:
+    #     image0 = image_0[:, :, ::-1]  # BGR to RGB
+    data = preprocess(image0, conf)
+    pred = model({"image": data["image"]})
+    pred["image_size"] = data["original_size"]
+    pred = {**pred, **data}
+    return pred
+
+
+@torch.no_grad()
+def main(
+    conf: Dict,
+    image_dir: Path,
+    export_dir: Optional[Path] = None,
+    as_half: bool = True,
+    image_list: Optional[Union[Path, List[str]]] = None,
+    feature_path: Optional[Path] = None,
+    overwrite: bool = False,
+) -> Path:
+    logger.info(
+        "Extracting local features with configuration:"
+        f"\n{pprint.pformat(conf)}"
+    )
+
+    dataset = ImageDataset(image_dir, conf["preprocessing"], image_list)
+    if feature_path is None:
+        feature_path = Path(export_dir, conf["output"] + ".h5")
+    feature_path.parent.mkdir(exist_ok=True, parents=True)
+    skip_names = set(
+        list_h5_names(feature_path)
+        if feature_path.exists() and not overwrite
+        else ()
+    )
+    dataset.names = [n for n in dataset.names if n not in skip_names]
+    if len(dataset.names) == 0:
+        logger.info("Skipping the extraction.")
+        return feature_path
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    Model = dynamic_load(extractors, conf["model"]["name"])
+    model = Model(conf["model"]).eval().to(device)
+
+    loader = torch.utils.data.DataLoader(
+        dataset, num_workers=1, shuffle=False, pin_memory=True
+    )
+    for idx, data in enumerate(tqdm(loader)):
+        name = dataset.names[idx]
+        pred = model({"image": data["image"].to(device, non_blocking=True)})
+        pred = {k: v[0].cpu().numpy() for k, v in pred.items()}
+
+        pred["image_size"] = original_size = data["original_size"][0].numpy()
+        if "keypoints" in pred:
+            size = np.array(data["image"].shape[-2:][::-1])
+            scales = (original_size / size).astype(np.float32)
+            pred["keypoints"] = (pred["keypoints"] + 0.5) * scales[None] - 0.5
+            if "scales" in pred:
+                pred["scales"] *= scales.mean()
+            # add keypoint uncertainties scaled to the original resolution
+            uncertainty = getattr(model, "detection_noise", 1) * scales.mean()
+
+        if as_half:
+            for k in pred:
+                dt = pred[k].dtype
+                if (dt == np.float32) and (dt != np.float16):
+                    pred[k] = pred[k].astype(np.float16)
+
+        with h5py.File(str(feature_path), "a", libver="latest") as fd:
+            try:
+                if name in fd:
+                    del fd[name]
+                grp = fd.create_group(name)
+                for k, v in pred.items():
+                    grp.create_dataset(k, data=v)
+                if "keypoints" in pred:
+                    grp["keypoints"].attrs["uncertainty"] = uncertainty
+            except OSError as error:
+                if "No space left on device" in error.args[0]:
+                    logger.error(
+                        "Out of disk space: storing features on disk can take "
+                        "significant space, did you enable the as_half flag?"
+                    )
+                    del grp, fd[name]
+                raise error
+
+        del pred
+
+    logger.info("Finished exporting features.")
+    return feature_path
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_dir", type=Path, required=True)
+    parser.add_argument("--export_dir", type=Path, required=True)
+    parser.add_argument(
+        "--conf",
+        type=str,
+        default="superpoint_aachen",
+        choices=list(confs.keys()),
+    )
+    parser.add_argument("--as_half", action="store_true")
+    parser.add_argument("--image_list", type=Path)
+    parser.add_argument("--feature_path", type=Path)
+    args = parser.parse_args()
+    main(confs[args.conf], args.image_dir, args.export_dir, args.as_half)

hloc/extractors/alike.py

@@ -0,0 +1,55 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+alike_path = Path(__file__).parent / "../../third_party/ALIKE"
+sys.path.append(str(alike_path))
+from alike import ALike as Alike_
+from alike import configs
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Alike(BaseModel):
+    default_conf = {
+        "model_name": "alike-t",  # 'alike-t', 'alike-s', 'alike-n', 'alike-l'
+        "use_relu": True,
+        "multiscale": False,
+        "max_keypoints": 1000,
+        "detection_threshold": 0.5,
+        "top_k": -1,
+        "sub_pixel": False,
+    }
+
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.net = Alike_(
+            **configs[conf["model_name"]],
+            device=device,
+            top_k=conf["top_k"],
+            scores_th=conf["detection_threshold"],
+            n_limit=conf["max_keypoints"],
+        )
+        logger.info("Load Alike model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        image = image.permute(0, 2, 3, 1).squeeze()
+        image = image.cpu().numpy() * 255.0
+        pred = self.net(image, sub_pixel=self.conf["sub_pixel"])
+
+        keypoints = pred["keypoints"]
+        descriptors = pred["descriptors"]
+        scores = pred["scores"]
+
+        return {
+            "keypoints": torch.from_numpy(keypoints)[None],
+            "scores": torch.from_numpy(scores)[None],
+            "descriptors": torch.from_numpy(descriptors.T)[None],
+        }

hloc/extractors/cosplace.py

@@ -0,0 +1,44 @@
+"""
+Code for loading models trained with CosPlace as a global features extractor
+for geolocalization through image retrieval.
+Multiple models are available with different backbones. Below is a summary of
+models available (backbone : list of available output descriptors
+dimensionality). For example you can use a model based on a ResNet50 with
+descriptors dimensionality 1024.
+    ResNet18:  [32, 64, 128, 256, 512]
+    ResNet50:  [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet101: [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet152: [32, 64, 128, 256, 512, 1024, 2048]
+    VGG16:     [    64, 128, 256, 512]
+
+CosPlace paper: https://arxiv.org/abs/2204.02287
+"""
+
+import torch
+import torchvision.transforms as tvf
+
+from ..utils.base_model import BaseModel
+
+
+class CosPlace(BaseModel):
+    default_conf = {"backbone": "ResNet50", "fc_output_dim": 2048}
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "gmberton/CosPlace",
+            "get_trained_model",
+            backbone=conf["backbone"],
+            fc_output_dim=conf["fc_output_dim"],
+        ).eval()
+
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        self.norm_rgb = tvf.Normalize(mean=mean, std=std)
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+        desc = self.net(image)
+        return {
+            "global_descriptor": desc,
+        }

hloc/extractors/d2net.py

@@ -0,0 +1,68 @@
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+d2net_path = Path(__file__).parent / "../../third_party/d2net"
+sys.path.append(str(d2net_path))
+from lib.model_test import D2Net as _D2Net
+from lib.pyramid import process_multiscale
+
+
+class D2Net(BaseModel):
+    default_conf = {
+        "model_name": "d2_tf.pth",
+        "checkpoint_dir": d2net_path / "models",
+        "use_relu": True,
+        "multiscale": False,
+        "max_keypoints": 1024,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        model_file = conf["checkpoint_dir"] / conf["model_name"]
+        if not model_file.exists():
+            model_file.parent.mkdir(exist_ok=True)
+            cmd = [
+                "wget",
+                "--quiet",
+                "https://dusmanu.com/files/d2-net/" + conf["model_name"],
+                "-O",
+                str(model_file),
+            ]
+            subprocess.run(cmd, check=True)
+
+        self.net = _D2Net(
+            model_file=model_file, use_relu=conf["use_relu"], use_cuda=False
+        )
+        logger.info("Load D2Net model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        image = image.flip(1)  # RGB -> BGR
+        norm = image.new_tensor([103.939, 116.779, 123.68])
+        image = image * 255 - norm.view(1, 3, 1, 1)  # caffe normalization
+
+        if self.conf["multiscale"]:
+            keypoints, scores, descriptors = process_multiscale(image, self.net)
+        else:
+            keypoints, scores, descriptors = process_multiscale(
+                image, self.net, scales=[1]
+            )
+        keypoints = keypoints[:, [1, 0]]  # (x, y) and remove the scale
+
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        keypoints = keypoints[idxs, :2]
+        descriptors = descriptors[idxs]
+        scores = scores[idxs]
+
+        return {
+            "keypoints": torch.from_numpy(keypoints)[None],
+            "scores": torch.from_numpy(scores)[None],
+            "descriptors": torch.from_numpy(descriptors.T)[None],
+        }

hloc/extractors/darkfeat.py

@@ -0,0 +1,64 @@
+import subprocess
+import sys
+from pathlib import Path
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+darkfeat_path = Path(__file__).parent / "../../third_party/DarkFeat"
+sys.path.append(str(darkfeat_path))
+from darkfeat import DarkFeat as DarkFeat_
+
+
+class DarkFeat(BaseModel):
+    default_conf = {
+        "model_name": "DarkFeat.pth",
+        "max_keypoints": 1000,
+        "detection_threshold": 0.5,
+        "sub_pixel": False,
+    }
+    weight_urls = {
+        "DarkFeat.pth": "https://drive.google.com/uc?id=1Thl6m8NcmQ7zSAF-1_xaFs3F4H8UU6HX&confirm=t",
+    }
+    proxy = "http://localhost:1080"
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        model_path = darkfeat_path / "checkpoints" / conf["model_name"]
+        link = self.weight_urls[conf["model_name"]]
+        if not model_path.exists():
+            model_path.parent.mkdir(exist_ok=True)
+            cmd_wo_proxy = ["gdown", link, "-O", str(model_path)]
+            cmd = ["gdown", link, "-O", str(model_path), "--proxy", self.proxy]
+            logger.info(
+                f"Downloading the DarkFeat model with `{cmd_wo_proxy}`."
+            )
+            try:
+                subprocess.run(cmd_wo_proxy, check=True)
+            except subprocess.CalledProcessError as e:
+                logger.info(f"Downloading the model failed `{e}`.")
+                logger.info(f"Downloading the DarkFeat model with `{cmd}`.")
+                try:
+                    subprocess.run(cmd, check=True)
+                except subprocess.CalledProcessError as e:
+                    logger.error("Failed to download the DarkFeat model.")
+                    raise e
+
+        self.net = DarkFeat_(model_path)
+        logger.info("Load DarkFeat model done.")
+
+    def _forward(self, data):
+        pred = self.net({"image": data["image"]})
+        keypoints = pred["keypoints"]
+        descriptors = pred["descriptors"]
+        scores = pred["scores"]
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        keypoints = keypoints[idxs, :2]
+        descriptors = descriptors[:, idxs]
+        scores = scores[idxs]
+        return {
+            "keypoints": keypoints[None],  # 1 x N x 2
+            "scores": scores[None],  # 1 x N
+            "descriptors": descriptors[None],  # 1 x 128 x N
+        }

hloc/extractors/dedode.py

@@ -0,0 +1,111 @@
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.transforms as transforms
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+dedode_path = Path(__file__).parent / "../../third_party/DeDoDe"
+sys.path.append(str(dedode_path))
+
+from DeDoDe import dedode_descriptor_B, dedode_detector_L
+from DeDoDe.utils import to_pixel_coords
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class DeDoDe(BaseModel):
+    default_conf = {
+        "name": "dedode",
+        "model_detector_name": "dedode_detector_L.pth",
+        "model_descriptor_name": "dedode_descriptor_B.pth",
+        "max_keypoints": 2000,
+        "match_threshold": 0.2,
+        "dense": False,  # Now fixed to be false
+    }
+    required_inputs = [
+        "image",
+    ]
+    weight_urls = {
+        "dedode_detector_L.pth": "https://github.com/Parskatt/DeDoDe/releases/download/dedode_pretrained_models/dedode_detector_L.pth",
+        "dedode_descriptor_B.pth": "https://github.com/Parskatt/DeDoDe/releases/download/dedode_pretrained_models/dedode_descriptor_B.pth",
+    }
+
+    # Initialize the line matcher
+    def _init(self, conf):
+        model_detector_path = (
+            dedode_path / "pretrained" / conf["model_detector_name"]
+        )
+        model_descriptor_path = (
+            dedode_path / "pretrained" / conf["model_descriptor_name"]
+        )
+
+        self.normalizer = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+        # Download the model.
+        if not model_detector_path.exists():
+            model_detector_path.parent.mkdir(exist_ok=True)
+            link = self.weight_urls[conf["model_detector_name"]]
+            cmd = ["wget", "--quiet", link, "-O", str(model_detector_path)]
+            logger.info(f"Downloading the DeDoDe detector model with `{cmd}`.")
+            subprocess.run(cmd, check=True)
+
+        if not model_descriptor_path.exists():
+            model_descriptor_path.parent.mkdir(exist_ok=True)
+            link = self.weight_urls[conf["model_descriptor_name"]]
+            cmd = ["wget", "--quiet", link, "-O", str(model_descriptor_path)]
+            logger.info(
+                f"Downloading the DeDoDe descriptor model with `{cmd}`."
+            )
+            subprocess.run(cmd, check=True)
+
+        # load the model
+        weights_detector = torch.load(model_detector_path, map_location="cpu")
+        weights_descriptor = torch.load(
+            model_descriptor_path, map_location="cpu"
+        )
+        self.detector = dedode_detector_L(
+            weights=weights_detector, device=device
+        )
+        self.descriptor = dedode_descriptor_B(
+            weights=weights_descriptor, device=device
+        )
+        logger.info("Load DeDoDe model done.")
+
+    def _forward(self, data):
+        """
+        data: dict, keys: {'image0','image1'}
+        image shape: N x C x H x W
+        color mode: RGB
+        """
+        img0 = self.normalizer(data["image"].squeeze()).float()[None]
+        H_A, W_A = img0.shape[2:]
+
+        # step 1: detect keypoints
+        detections_A = None
+        batch_A = {"image": img0}
+        if self.conf["dense"]:
+            detections_A = self.detector.detect_dense(batch_A)
+        else:
+            detections_A = self.detector.detect(
+                batch_A, num_keypoints=self.conf["max_keypoints"]
+            )
+        keypoints_A, P_A = detections_A["keypoints"], detections_A["confidence"]
+
+        # step 2: describe keypoints
+        # dim: 1 x N x 256
+        description_A = self.descriptor.describe_keypoints(
+            batch_A, keypoints_A
+        )["descriptions"]
+        keypoints_A = to_pixel_coords(keypoints_A, H_A, W_A)
+
+        return {
+            "keypoints": keypoints_A,  # 1 x N x 2
+            "descriptors": description_A.permute(0, 2, 1),  # 1 x 256 x N
+            "scores": P_A,  # 1 x N
+        }

hloc/extractors/dir.py

@@ -0,0 +1,81 @@
+import os
+import sys
+from pathlib import Path
+from zipfile import ZipFile
+
+import gdown
+import sklearn
+import torch
+
+from ..utils.base_model import BaseModel
+
+sys.path.append(
+    str(Path(__file__).parent / "../../third_party/deep-image-retrieval")
+)
+os.environ["DB_ROOT"] = ""  # required by dirtorch
+
+from dirtorch.extract_features import load_model  # noqa: E402
+from dirtorch.utils import common  # noqa: E402
+
+# The DIR model checkpoints (pickle files) include sklearn.decomposition.pca,
+# which has been deprecated in sklearn v0.24
+# and must be explicitly imported with `from sklearn.decomposition import PCA`.
+# This is a hacky workaround to maintain forward compatibility.
+sys.modules["sklearn.decomposition.pca"] = sklearn.decomposition._pca
+
+
+class DIR(BaseModel):
+    default_conf = {
+        "model_name": "Resnet-101-AP-GeM",
+        "whiten_name": "Landmarks_clean",
+        "whiten_params": {
+            "whitenp": 0.25,
+            "whitenv": None,
+            "whitenm": 1.0,
+        },
+        "pooling": "gem",
+        "gemp": 3,
+    }
+    required_inputs = ["image"]
+
+    dir_models = {
+        "Resnet-101-AP-GeM": "https://docs.google.com/uc?export=download&id=1UWJGDuHtzaQdFhSMojoYVQjmCXhIwVvy",
+    }
+
+    def _init(self, conf):
+        checkpoint = Path(
+            torch.hub.get_dir(), "dirtorch", conf["model_name"] + ".pt"
+        )
+        if not checkpoint.exists():
+            checkpoint.parent.mkdir(exist_ok=True, parents=True)
+            link = self.dir_models[conf["model_name"]]
+            gdown.download(str(link), str(checkpoint) + ".zip", quiet=False)
+            zf = ZipFile(str(checkpoint) + ".zip", "r")
+            zf.extractall(checkpoint.parent)
+            zf.close()
+            os.remove(str(checkpoint) + ".zip")
+
+        self.net = load_model(checkpoint, False)  # first load on CPU
+        if conf["whiten_name"]:
+            assert conf["whiten_name"] in self.net.pca
+
+    def _forward(self, data):
+        image = data["image"]
+        assert image.shape[1] == 3
+        mean = self.net.preprocess["mean"]
+        std = self.net.preprocess["std"]
+        image = image - image.new_tensor(mean)[:, None, None]
+        image = image / image.new_tensor(std)[:, None, None]
+
+        desc = self.net(image)
+        desc = desc.unsqueeze(0)  # batch dimension
+        if self.conf["whiten_name"]:
+            pca = self.net.pca[self.conf["whiten_name"]]
+            desc = common.whiten_features(
+                desc.cpu().numpy(), pca, **self.conf["whiten_params"]
+            )
+            desc = torch.from_numpy(desc)
+
+        return {
+            "global_descriptor": desc,
+        }

hloc/extractors/disk.py

@@ -0,0 +1,35 @@
+import kornia
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+
+class DISK(BaseModel):
+    default_conf = {
+        "weights": "depth",
+        "max_keypoints": None,
+        "nms_window_size": 5,
+        "detection_threshold": 0.0,
+        "pad_if_not_divisible": True,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.model = kornia.feature.DISK.from_pretrained(conf["weights"])
+        logger.info("Load DISK model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        features = self.model(
+            image,
+            n=self.conf["max_keypoints"],
+            window_size=self.conf["nms_window_size"],
+            score_threshold=self.conf["detection_threshold"],
+            pad_if_not_divisible=self.conf["pad_if_not_divisible"],
+        )
+        return {
+            "keypoints": [f.keypoints for f in features][0][None],
+            "scores": [f.detection_scores for f in features][0][None],
+            "descriptors": [f.descriptors.t() for f in features][0][None],
+        }

hloc/extractors/dog.py

@@ -0,0 +1,135 @@
+import kornia
+import numpy as np
+import pycolmap
+import torch
+from kornia.feature.laf import (
+    extract_patches_from_pyramid,
+    laf_from_center_scale_ori,
+)
+
+from ..utils.base_model import BaseModel
+
+EPS = 1e-6
+
+
+def sift_to_rootsift(x):
+    x = x / (np.linalg.norm(x, ord=1, axis=-1, keepdims=True) + EPS)
+    x = np.sqrt(x.clip(min=EPS))
+    x = x / (np.linalg.norm(x, axis=-1, keepdims=True) + EPS)
+    return x
+
+
+class DoG(BaseModel):
+    default_conf = {
+        "options": {
+            "first_octave": 0,
+            "peak_threshold": 0.01,
+        },
+        "descriptor": "rootsift",
+        "max_keypoints": -1,
+        "patch_size": 32,
+        "mr_size": 12,
+    }
+    required_inputs = ["image"]
+    detection_noise = 1.0
+    max_batch_size = 1024
+
+    def _init(self, conf):
+        if conf["descriptor"] == "sosnet":
+            self.describe = kornia.feature.SOSNet(pretrained=True)
+        elif conf["descriptor"] == "hardnet":
+            self.describe = kornia.feature.HardNet(pretrained=True)
+        elif conf["descriptor"] not in ["sift", "rootsift"]:
+            raise ValueError(f'Unknown descriptor: {conf["descriptor"]}')
+
+        self.sift = None  # lazily instantiated on the first image
+        self.dummy_param = torch.nn.Parameter(torch.empty(0))
+        self.device = torch.device("cpu")
+
+    def to(self, *args, **kwargs):
+        device = kwargs.get("device")
+        if device is None:
+            match = [a for a in args if isinstance(a, (torch.device, str))]
+            if len(match) > 0:
+                device = match[0]
+        if device is not None:
+            self.device = torch.device(device)
+        return super().to(*args, **kwargs)
+
+    def _forward(self, data):
+        image = data["image"]
+        image_np = image.cpu().numpy()[0, 0]
+        assert image.shape[1] == 1
+        assert image_np.min() >= -EPS and image_np.max() <= 1 + EPS
+
+        if self.sift is None:
+            device = self.dummy_param.device
+            use_gpu = pycolmap.has_cuda and device.type == "cuda"
+            options = {**self.conf["options"]}
+            if self.conf["descriptor"] == "rootsift":
+                options["normalization"] = pycolmap.Normalization.L1_ROOT
+            else:
+                options["normalization"] = pycolmap.Normalization.L2
+            self.sift = pycolmap.Sift(
+                options=pycolmap.SiftExtractionOptions(options),
+                device=getattr(pycolmap.Device, "cuda" if use_gpu else "cpu"),
+            )
+        keypoints, descriptors = self.sift.extract(image_np)
+        scales = keypoints[:, 2]
+        oris = np.rad2deg(keypoints[:, 3])
+
+        if self.conf["descriptor"] in ["sift", "rootsift"]:
+            # We still renormalize because COLMAP does not normalize well,
+            # maybe due to numerical errors
+            if self.conf["descriptor"] == "rootsift":
+                descriptors = sift_to_rootsift(descriptors)
+            descriptors = torch.from_numpy(descriptors)
+        elif self.conf["descriptor"] in ("sosnet", "hardnet"):
+            center = keypoints[:, :2] + 0.5
+            laf_scale = scales * self.conf["mr_size"] / 2
+            laf_ori = -oris
+            lafs = laf_from_center_scale_ori(
+                torch.from_numpy(center)[None],
+                torch.from_numpy(laf_scale)[None, :, None, None],
+                torch.from_numpy(laf_ori)[None, :, None],
+            ).to(image.device)
+            patches = extract_patches_from_pyramid(
+                image, lafs, PS=self.conf["patch_size"]
+            )[0]
+            descriptors = patches.new_zeros((len(patches), 128))
+            if len(patches) > 0:
+                for start_idx in range(0, len(patches), self.max_batch_size):
+                    end_idx = min(len(patches), start_idx + self.max_batch_size)
+                    descriptors[start_idx:end_idx] = self.describe(
+                        patches[start_idx:end_idx]
+                    )
+        else:
+            raise ValueError(f'Unknown descriptor: {self.conf["descriptor"]}')
+
+        keypoints = torch.from_numpy(keypoints[:, :2])  # keep only x, y
+        scales = torch.from_numpy(scales)
+        oris = torch.from_numpy(oris)
+        scores = keypoints.new_zeros(len(keypoints))  # no scores for SIFT yet
+
+        if self.conf["max_keypoints"] != -1:
+            # TODO: check that the scores from PyCOLMAP are 100% correct,
+            # follow https://github.com/mihaidusmanu/pycolmap/issues/8
+            max_number = (
+                scores.shape[0]
+                if scores.shape[0] < self.conf["max_keypoints"]
+                else self.conf["max_keypoints"]
+            )
+            values, indices = torch.topk(scores, max_number)
+            keypoints = keypoints[indices]
+            scales = scales[indices]
+            oris = oris[indices]
+            scores = scores[indices]
+            descriptors = descriptors[indices]
+
+        return {
+            "keypoints": keypoints[None],
+            "scales": scales[None],
+            "oris": oris[None],
+            "scores": scores[None],
+            "descriptors": descriptors.T[None],
+        }

hloc/extractors/eigenplaces.py

@@ -0,0 +1,57 @@
+"""
+Code for loading models trained with EigenPlaces (or CosPlace) as a global
+features extractor for geolocalization through image retrieval.
+Multiple models are available with different backbones. Below is a summary of
+models available (backbone : list of available output descriptors
+dimensionality). For example you can use a model based on a ResNet50 with
+descriptors dimensionality 1024.
+
+EigenPlaces trained models:
+    ResNet18:  [     256, 512]
+    ResNet50:  [128, 256, 512, 2048]
+    ResNet101: [128, 256, 512, 2048]
+    VGG16:     [     512]
+
+CosPlace trained models:
+    ResNet18:  [32, 64, 128, 256, 512]
+    ResNet50:  [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet101: [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet152: [32, 64, 128, 256, 512, 1024, 2048]
+    VGG16:     [    64, 128, 256, 512]
+
+EigenPlaces paper (ICCV 2023): https://arxiv.org/abs/2308.10832
+CosPlace paper (CVPR 2022): https://arxiv.org/abs/2204.02287
+"""
+
+import torch
+import torchvision.transforms as tvf
+
+from ..utils.base_model import BaseModel
+
+
+class EigenPlaces(BaseModel):
+    default_conf = {
+        "variant": "EigenPlaces",
+        "backbone": "ResNet101",
+        "fc_output_dim": 2048,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "gmberton/" + conf["variant"],
+            "get_trained_model",
+            backbone=conf["backbone"],
+            fc_output_dim=conf["fc_output_dim"],
+        ).eval()
+
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        self.norm_rgb = tvf.Normalize(mean=mean, std=std)
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+        desc = self.net(image)
+        return {
+            "global_descriptor": desc,
+        }

hloc/extractors/example.py

@@ -0,0 +1,56 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import logger
+from ..utils.base_model import BaseModel
+
+example_path = Path(__file__).parent / "../../third_party/example"
+sys.path.append(str(example_path))
+
+# import some modules here
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Example(BaseModel):
+    # change to your default configs
+    default_conf = {
+        "name": "example",
+        "keypoint_threshold": 0.1,
+        "max_keypoints": 2000,
+        "model_name": "model.pth",
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        # set checkpoints paths if needed
+        model_path = example_path / "checkpoints" / f'{conf["model_name"]}'
+        if not model_path.exists():
+            logger.info(f"No model found at {model_path}")
+
+        # init model
+        self.net = callable
+        # self.net = ExampleNet(is_test=True)
+        state_dict = torch.load(model_path, map_location="cpu")
+        self.net.load_state_dict(state_dict["model_state"])
+        logger.info("Load example model done.")
+
+    def _forward(self, data):
+        # data: dict, keys: 'image'
+        # image color mode: RGB
+        # image value range in [0, 1]
+        image = data["image"]
+
+        # B: batch size, N: number of keypoints
+        # keypoints shape: B x N x 2, type: torch tensor
+        # scores shape: B x N, type: torch tensor
+        # descriptors shape: B x 128 x N, type: torch tensor
+        keypoints, scores, descriptors = self.net(image)
+
+        return {
+            "keypoints": keypoints,
+            "scores": scores,
+            "descriptors": descriptors,
+        }

hloc/extractors/fire.py

@@ -0,0 +1,72 @@
+import logging
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.transforms as tvf
+
+from ..utils.base_model import BaseModel
+
+logger = logging.getLogger(__name__)
+fire_path = Path(__file__).parent / "../../third_party/fire"
+sys.path.append(str(fire_path))
+
+
+import fire_network
+
+
+class FIRe(BaseModel):
+    default_conf = {
+        "global": True,
+        "asmk": False,
+        "model_name": "fire_SfM_120k.pth",
+        "scales": [2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25],  # default params
+        "features_num": 1000,  # TODO:not supported now
+        "asmk_name": "asmk_codebook.bin",  # TODO:not supported now
+        "config_name": "eval_fire.yml",
+    }
+    required_inputs = ["image"]
+
+    # Models exported using
+    fire_models = {
+        "fire_SfM_120k.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/official/fire.pth",
+        "fire_imagenet.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/pretraining/fire_imagenet.pth",
+    }
+
+    def _init(self, conf):
+        assert conf["model_name"] in self.fire_models.keys()
+        # Config paths
+        model_path = fire_path / "model" / conf["model_name"]
+
+        # Download the model.
+        if not model_path.exists():
+            model_path.parent.mkdir(exist_ok=True)
+            link = self.fire_models[conf["model_name"]]
+            cmd = ["wget", "--quiet", link, "-O", str(model_path)]
+            logger.info(f"Downloading the FIRe model with `{cmd}`.")
+            subprocess.run(cmd, check=True)
+
+        logger.info("Loading fire model...")
+
+        # Load net
+        state = torch.load(model_path)
+        state["net_params"]["pretrained"] = None
+        net = fire_network.init_network(**state["net_params"])
+        net.load_state_dict(state["state_dict"])
+        self.net = net
+
+        self.norm_rgb = tvf.Normalize(
+            **dict(zip(["mean", "std"], net.runtime["mean_std"]))
+        )
+
+        # params
+        self.scales = conf["scales"]
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+
+        # Feature extraction.
+        desc = self.net.forward_global(image, scales=self.scales)
+
+        return {"global_descriptor": desc}

hloc/extractors/fire_local.py

@@ -0,0 +1,84 @@
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.transforms as tvf
+
+from .. import logger
+from ..utils.base_model import BaseModel
+
+fire_path = Path(__file__).parent / "../../third_party/fire"
+
+sys.path.append(str(fire_path))
+
+
+import fire_network
+
+EPS = 1e-6
+
+
+class FIRe(BaseModel):
+    default_conf = {
+        "global": True,
+        "asmk": False,
+        "model_name": "fire_SfM_120k.pth",
+        "scales": [2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25],  # default params
+        "features_num": 1000,
+        "asmk_name": "asmk_codebook.bin",
+        "config_name": "eval_fire.yml",
+    }
+    required_inputs = ["image"]
+
+    # Models exported using
+    fire_models = {
+        "fire_SfM_120k.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/official/fire.pth",
+        "fire_imagenet.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/pretraining/fire_imagenet.pth",
+    }
+
+    def _init(self, conf):
+        assert conf["model_name"] in self.fire_models.keys()
+
+        # Config paths
+        model_path = fire_path / "model" / conf["model_name"]
+        config_path = fire_path / conf["config_name"]  # noqa: F841
+        asmk_bin_path = fire_path / "model" / conf["asmk_name"]  # noqa: F841
+
+        # Download the model.
+        if not model_path.exists():
+            model_path.parent.mkdir(exist_ok=True)
+            link = self.fire_models[conf["model_name"]]
+            cmd = ["wget", "--quiet", link, "-O", str(model_path)]
+            logger.info(f"Downloading the FIRe model with `{cmd}`.")
+            subprocess.run(cmd, check=True)
+
+        logger.info("Loading fire model...")
+
+        # Load net
+        state = torch.load(model_path)
+        state["net_params"]["pretrained"] = None
+        net = fire_network.init_network(**state["net_params"])
+        net.load_state_dict(state["state_dict"])
+        self.net = net
+
+        self.norm_rgb = tvf.Normalize(
+            **dict(zip(["mean", "std"], net.runtime["mean_std"]))
+        )
+
+        # params
+        self.scales = conf["scales"]
+        self.features_num = conf["features_num"]
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+
+        local_desc = self.net.forward_local(
+            image, features_num=self.features_num, scales=self.scales
+        )
+
+        logger.info(f"output[0].shape = {local_desc[0].shape}\n")
+
+        return {
+            # 'global_descriptor': desc
+            "local_descriptor": local_desc
+        }

hloc/extractors/lanet.py

@@ -0,0 +1,66 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+lib_path = Path(__file__).parent / "../../third_party"
+sys.path.append(str(lib_path))
+from lanet.network_v0.model import PointModel
+
+lanet_path = Path(__file__).parent / "../../third_party/lanet"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class LANet(BaseModel):
+    default_conf = {
+        "model_name": "v0",
+        "keypoint_threshold": 0.1,
+        "max_keypoints": 1024,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        model_path = (
+            lanet_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        )
+        if not model_path.exists():
+            logger.warning(f"No model found at {model_path}, start downloading")
+        self.net = PointModel(is_test=True)
+        state_dict = torch.load(model_path, map_location="cpu")
+        self.net.load_state_dict(state_dict["model_state"])
+        logger.info("Load LANet model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        keypoints, scores, descriptors = self.net(image)
+        _, _, Hc, Wc = descriptors.shape
+
+        # Scores & Descriptors
+        kpts_score = torch.cat([keypoints, scores], dim=1).view(3, -1).t()
+        descriptors = descriptors.view(256, Hc, Wc).view(256, -1).t()
+
+        # Filter based on confidence threshold
+        descriptors = descriptors[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
+        kpts_score = kpts_score[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
+        keypoints = kpts_score[:, 1:]
+        scores = kpts_score[:, 0]
+
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        keypoints = keypoints[idxs, :2]
+        descriptors = descriptors[idxs]
+        scores = scores[idxs]
+
+        return {
+            "keypoints": keypoints[None],
+            "scores": scores[None],
+            "descriptors": descriptors.T[None],
+        }

hloc/extractors/netvlad.py

@@ -0,0 +1,152 @@
+import subprocess
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from scipy.io import loadmat
+
+from .. import logger
+from ..utils.base_model import BaseModel
+
+EPS = 1e-6
+
+
+class NetVLADLayer(nn.Module):
+    def __init__(self, input_dim=512, K=64, score_bias=False, intranorm=True):
+        super().__init__()
+        self.score_proj = nn.Conv1d(
+            input_dim, K, kernel_size=1, bias=score_bias
+        )
+        centers = nn.parameter.Parameter(torch.empty([input_dim, K]))
+        nn.init.xavier_uniform_(centers)
+        self.register_parameter("centers", centers)
+        self.intranorm = intranorm
+        self.output_dim = input_dim * K
+
+    def forward(self, x):
+        b = x.size(0)
+        scores = self.score_proj(x)
+        scores = F.softmax(scores, dim=1)
+        diff = x.unsqueeze(2) - self.centers.unsqueeze(0).unsqueeze(-1)
+        desc = (scores.unsqueeze(1) * diff).sum(dim=-1)
+        if self.intranorm:
+            # From the official MATLAB implementation.
+            desc = F.normalize(desc, dim=1)
+        desc = desc.view(b, -1)
+        desc = F.normalize(desc, dim=1)
+        return desc
+
+
+class NetVLAD(BaseModel):
+    default_conf = {"model_name": "VGG16-NetVLAD-Pitts30K", "whiten": True}
+    required_inputs = ["image"]
+
+    # Models exported using
+    # https://github.com/uzh-rpg/netvlad_tf_open/blob/master/matlab/net_class2struct.m.
+    dir_models = {
+        "VGG16-NetVLAD-Pitts30K": "https://cvg-data.inf.ethz.ch/hloc/netvlad/Pitts30K_struct.mat",
+        "VGG16-NetVLAD-TokyoTM": "https://cvg-data.inf.ethz.ch/hloc/netvlad/TokyoTM_struct.mat",
+    }
+
+    def _init(self, conf):
+        assert conf["model_name"] in self.dir_models.keys()
+
+        # Download the checkpoint.
+        checkpoint = Path(
+            torch.hub.get_dir(), "netvlad", conf["model_name"] + ".mat"
+        )
+        if not checkpoint.exists():
+            checkpoint.parent.mkdir(exist_ok=True, parents=True)
+            link = self.dir_models[conf["model_name"]]
+            cmd = ["wget", "--quiet", link, "-O", str(checkpoint)]
+            logger.info(f"Downloading the NetVLAD model with `{cmd}`.")
+            subprocess.run(cmd, check=True)
+
+        # Create the network.
+        # Remove classification head.
+        backbone = list(models.vgg16().children())[0]
+        # Remove last ReLU + MaxPool2d.
+        self.backbone = nn.Sequential(*list(backbone.children())[:-2])
+
+        self.netvlad = NetVLADLayer()
+
+        if conf["whiten"]:
+            self.whiten = nn.Linear(self.netvlad.output_dim, 4096)
+
+        # Parse MATLAB weights using https://github.com/uzh-rpg/netvlad_tf_open
+        mat = loadmat(checkpoint, struct_as_record=False, squeeze_me=True)
+
+        # CNN weights.
+        for layer, mat_layer in zip(
+            self.backbone.children(), mat["net"].layers
+        ):
+            if isinstance(layer, nn.Conv2d):
+                w = mat_layer.weights[0]  # Shape: S x S x IN x OUT
+                b = mat_layer.weights[1]  # Shape: OUT
+                # Prepare for PyTorch - enforce float32 and right shape.
+                # w should have shape: OUT x IN x S x S
+                # b should have shape: OUT
+                w = torch.tensor(w).float().permute([3, 2, 0, 1])
+                b = torch.tensor(b).float()
+                # Update layer weights.
+                layer.weight = nn.Parameter(w)
+                layer.bias = nn.Parameter(b)
+
+        # NetVLAD weights.
+        score_w = mat["net"].layers[30].weights[0]  # D x K
+        # centers are stored as opposite in official MATLAB code
+        center_w = -mat["net"].layers[30].weights[1]  # D x K
+        # Prepare for PyTorch - make sure it is float32 and has right shape.
+        # score_w should have shape K x D x 1
+        # center_w should have shape D x K
+        score_w = torch.tensor(score_w).float().permute([1, 0]).unsqueeze(-1)
+        center_w = torch.tensor(center_w).float()
+        # Update layer weights.
+        self.netvlad.score_proj.weight = nn.Parameter(score_w)
+        self.netvlad.centers = nn.Parameter(center_w)
+
+        # Whitening weights.
+        if conf["whiten"]:
+            w = mat["net"].layers[33].weights[0]  # Shape: 1 x 1 x IN x OUT
+            b = mat["net"].layers[33].weights[1]  # Shape: OUT
+            # Prepare for PyTorch - make sure it is float32 and has right shape
+            w = torch.tensor(w).float().squeeze().permute([1, 0])  # OUT x IN
+            b = torch.tensor(b.squeeze()).float()  # Shape: OUT
+            # Update layer weights.
+            self.whiten.weight = nn.Parameter(w)
+            self.whiten.bias = nn.Parameter(b)
+
+        # Preprocessing parameters.
+        self.preprocess = {
+            "mean": mat["net"].meta.normalization.averageImage[0, 0],
+            "std": np.array([1, 1, 1], dtype=np.float32),
+        }
+
+    def _forward(self, data):
+        image = data["image"]
+        assert image.shape[1] == 3
+        assert image.min() >= -EPS and image.max() <= 1 + EPS
+        image = torch.clamp(image * 255, 0.0, 255.0)  # Input should be 0-255.
+        mean = self.preprocess["mean"]
+        std = self.preprocess["std"]
+        image = image - image.new_tensor(mean).view(1, -1, 1, 1)
+        image = image / image.new_tensor(std).view(1, -1, 1, 1)
+
+        # Feature extraction.
+        descriptors = self.backbone(image)
+        b, c, _, _ = descriptors.size()
+        descriptors = descriptors.view(b, c, -1)
+
+        # NetVLAD layer.
+        descriptors = F.normalize(descriptors, dim=1)  # Pre-normalization.
+        desc = self.netvlad(descriptors)
+
+        # Whiten if needed.
+        if hasattr(self, "whiten"):
+            desc = self.whiten(desc)
+            desc = F.normalize(desc, dim=1)  # Final L2 normalization.
+
+        return {"global_descriptor": desc}

hloc/extractors/openibl.py

@@ -0,0 +1,26 @@
+import torch
+import torchvision.transforms as tvf
+
+from ..utils.base_model import BaseModel
+
+
+class OpenIBL(BaseModel):
+    default_conf = {
+        "model_name": "vgg16_netvlad",
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "yxgeee/OpenIBL", conf["model_name"], pretrained=True
+        ).eval()
+        mean = [0.48501960784313836, 0.4579568627450961, 0.4076039215686255]
+        std = [0.00392156862745098, 0.00392156862745098, 0.00392156862745098]
+        self.norm_rgb = tvf.Normalize(mean=mean, std=std)
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+        desc = self.net(image)
+        return {
+            "global_descriptor": desc,
+        }

hloc/extractors/r2d2.py

@@ -0,0 +1,65 @@
+import sys
+from pathlib import Path
+
+import torchvision.transforms as tvf
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+r2d2_path = Path(__file__).parent / "../../third_party/r2d2"
+sys.path.append(str(r2d2_path))
+from extract import NonMaxSuppression, extract_multiscale, load_network
+
+
+class R2D2(BaseModel):
+    default_conf = {
+        "model_name": "r2d2_WASF_N16.pt",
+        "max_keypoints": 5000,
+        "scale_factor": 2**0.25,
+        "min_size": 256,
+        "max_size": 1024,
+        "min_scale": 0,
+        "max_scale": 1,
+        "reliability_threshold": 0.7,
+        "repetability_threshold": 0.7,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        model_fn = r2d2_path / "models" / conf["model_name"]
+        self.norm_rgb = tvf.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+        self.net = load_network(model_fn)
+        self.detector = NonMaxSuppression(
+            rel_thr=conf["reliability_threshold"],
+            rep_thr=conf["repetability_threshold"],
+        )
+        logger.info("Load R2D2 model done.")
+
+    def _forward(self, data):
+        img = data["image"]
+        img = self.norm_rgb(img)
+
+        xys, desc, scores = extract_multiscale(
+            self.net,
+            img,
+            self.detector,
+            scale_f=self.conf["scale_factor"],
+            min_size=self.conf["min_size"],
+            max_size=self.conf["max_size"],
+            min_scale=self.conf["min_scale"],
+            max_scale=self.conf["max_scale"],
+        )
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        xy = xys[idxs, :2]
+        desc = desc[idxs].t()
+        scores = scores[idxs]
+
+        pred = {
+            "keypoints": xy[None],
+            "descriptors": desc[None],
+            "scores": scores[None],
+        }
+        return pred

hloc/extractors/rekd.py

@@ -0,0 +1,72 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+rekd_path = Path(__file__).parent / "../../third_party"
+sys.path.append(str(rekd_path))
+from REKD.training.model.REKD import REKD as REKD_
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class REKD(BaseModel):
+    default_conf = {
+        "model_name": "v0",
+        "keypoint_threshold": 0.1,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        model_path = (
+            rekd_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        )
+        if not model_path.exists():
+            print(f"No model found at {model_path}")
+        self.net = REKD_(is_test=True)
+        state_dict = torch.load(model_path, map_location="cpu")
+        self.net.load_state_dict(state_dict["model_state"])
+        logger.info("Load REKD model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        keypoints, scores, descriptors = self.net(image)
+        _, _, Hc, Wc = descriptors.shape
+
+        # Scores & Descriptors
+        kpts_score = (
+            torch.cat([keypoints, scores], dim=1)
+            .view(3, -1)
+            .t()
+            .cpu()
+            .detach()
+            .numpy()
+        )
+        descriptors = (
+            descriptors.view(256, Hc, Wc)
+            .view(256, -1)
+            .t()
+            .cpu()
+            .detach()
+            .numpy()
+        )
+
+        # Filter based on confidence threshold
+        descriptors = descriptors[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
+        kpts_score = kpts_score[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
+        keypoints = kpts_score[:, 1:]
+        scores = kpts_score[:, 0]
+
+        return {
+            "keypoints": torch.from_numpy(keypoints)[None],
+            "scores": torch.from_numpy(scores)[None],
+            "descriptors": torch.from_numpy(descriptors.T)[None],
+        }

hloc/extractors/rord.py

@@ -0,0 +1,76 @@
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+rord_path = Path(__file__).parent / "../../third_party"
+sys.path.append(str(rord_path))
+from RoRD.lib.model_test import D2Net as _RoRD
+from RoRD.lib.pyramid import process_multiscale
+
+
+class RoRD(BaseModel):
+    default_conf = {
+        "model_name": "rord.pth",
+        "checkpoint_dir": rord_path / "RoRD" / "models",
+        "use_relu": True,
+        "multiscale": False,
+        "max_keypoints": 1024,
+    }
+    required_inputs = ["image"]
+    weight_urls = {
+        "rord.pth": "https://drive.google.com/uc?id=12414ZGKwgPAjNTGtNrlB4VV9l7W76B2o&confirm=t",
+    }
+    proxy = "http://localhost:1080"
+
+    def _init(self, conf):
+        model_path = conf["checkpoint_dir"] / conf["model_name"]
+        link = self.weight_urls[conf["model_name"]]
+        if not model_path.exists():
+            model_path.parent.mkdir(exist_ok=True)
+            cmd_wo_proxy = ["gdown", link, "-O", str(model_path)]
+            cmd = ["gdown", link, "-O", str(model_path), "--proxy", self.proxy]
+            logger.info(f"Downloading the RoRD model with `{cmd_wo_proxy}`.")
+            try:
+                subprocess.run(cmd_wo_proxy, check=True)
+            except subprocess.CalledProcessError as e:
+                logger.info(f"Downloading failed {e}.")
+                logger.info(f"Downloading the RoRD model with {cmd}.")
+                try:
+                    subprocess.run(cmd, check=True)
+                except subprocess.CalledProcessError as e:
+                    logger.error(f"Failed to download the RoRD model: {e}")
+        self.net = _RoRD(
+            model_file=model_path, use_relu=conf["use_relu"], use_cuda=False
+        )
+        logger.info("Load RoRD model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        image = image.flip(1)  # RGB -> BGR
+        norm = image.new_tensor([103.939, 116.779, 123.68])
+        image = image * 255 - norm.view(1, 3, 1, 1)  # caffe normalization
+
+        if self.conf["multiscale"]:
+            keypoints, scores, descriptors = process_multiscale(image, self.net)
+        else:
+            keypoints, scores, descriptors = process_multiscale(
+                image, self.net, scales=[1]
+            )
+        keypoints = keypoints[:, [1, 0]]  # (x, y) and remove the scale
+
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        keypoints = keypoints[idxs, :2]
+        descriptors = descriptors[idxs]
+        scores = scores[idxs]
+
+        return {
+            "keypoints": torch.from_numpy(keypoints)[None],
+            "scores": torch.from_numpy(scores)[None],
+            "descriptors": torch.from_numpy(descriptors.T)[None],
+        }

hloc/extractors/sfd2.py

@@ -0,0 +1,41 @@
+import sys
+from pathlib import Path
+
+import torchvision.transforms as tvf
+
+from .. import logger
+from ..utils.base_model import BaseModel
+
+tp_path = Path(__file__).parent / "../../third_party"
+sys.path.append(str(tp_path))
+from pram.nets.sfd2 import load_sfd2
+
+
+class SFD2(BaseModel):
+    default_conf = {
+        "max_keypoints": 4096,
+        "model_name": "sfd2_20230511_210205_resnet4x.79.pth",
+        "conf_th": 0.001,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.conf = {**self.default_conf, **conf}
+        self.norm_rgb = tvf.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+        model_path = tp_path / "pram" / "weights" / self.conf["model_name"]
+        self.net = load_sfd2(weight_path=model_path).eval()
+
+        logger.info("Load SFD2 model done.")
+
+    def _forward(self, data):
+        pred = self.net.extract_local_global(
+            data={"image": self.norm_rgb(data["image"])}, config=self.conf
+        )
+        out = {
+            "keypoints": pred["keypoints"][0][None],
+            "scores": pred["scores"][0][None],
+            "descriptors": pred["descriptors"][0][None],
+        }
+        return out

hloc/extractors/sift.py

@@ -0,0 +1,225 @@
+import warnings
+
+import cv2
+import numpy as np
+import torch
+from kornia.color import rgb_to_grayscale
+from omegaconf import OmegaConf
+from packaging import version
+
+try:
+    import pycolmap
+except ImportError:
+    pycolmap = None
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+
+def filter_dog_point(
+    points, scales, angles, image_shape, nms_radius, scores=None
+):
+    h, w = image_shape
+    ij = np.round(points - 0.5).astype(int).T[::-1]
+
+    # Remove duplicate points (identical coordinates).
+    # Pick highest scale or score
+    s = scales if scores is None else scores
+    buffer = np.zeros((h, w))
+    np.maximum.at(buffer, tuple(ij), s)
+    keep = np.where(buffer[tuple(ij)] == s)[0]
+
+    # Pick lowest angle (arbitrary).
+    ij = ij[:, keep]
+    buffer[:] = np.inf
+    o_abs = np.abs(angles[keep])
+    np.minimum.at(buffer, tuple(ij), o_abs)
+    mask = buffer[tuple(ij)] == o_abs
+    ij = ij[:, mask]
+    keep = keep[mask]
+
+    if nms_radius > 0:
+        # Apply NMS on the remaining points
+        buffer[:] = 0
+        buffer[tuple(ij)] = s[keep]  # scores or scale
+
+        local_max = torch.nn.functional.max_pool2d(
+            torch.from_numpy(buffer).unsqueeze(0),
+            kernel_size=nms_radius * 2 + 1,
+            stride=1,
+            padding=nms_radius,
+        ).squeeze(0)
+        is_local_max = buffer == local_max.numpy()
+        keep = keep[is_local_max[tuple(ij)]]
+    return keep
+
+
+def sift_to_rootsift(x: torch.Tensor, eps=1e-6) -> torch.Tensor:
+    x = torch.nn.functional.normalize(x, p=1, dim=-1, eps=eps)
+    x.clip_(min=eps).sqrt_()
+    return torch.nn.functional.normalize(x, p=2, dim=-1, eps=eps)
+
+
+def run_opencv_sift(features: cv2.Feature2D, image: np.ndarray) -> np.ndarray:
+    """
+    Detect keypoints using OpenCV Detector.
+    Optionally, perform description.
+    Args:
+        features: OpenCV based keypoints detector and descriptor
+        image: Grayscale image of uint8 data type
+    Returns:
+        keypoints: 1D array of detected cv2.KeyPoint
+        scores: 1D array of responses
+        descriptors: 1D array of descriptors
+    """
+    detections, descriptors = features.detectAndCompute(image, None)
+    points = np.array([k.pt for k in detections], dtype=np.float32)
+    scores = np.array([k.response for k in detections], dtype=np.float32)
+    scales = np.array([k.size for k in detections], dtype=np.float32)
+    angles = np.deg2rad(
+        np.array([k.angle for k in detections], dtype=np.float32)
+    )
+    return points, scores, scales, angles, descriptors
+
+
+class SIFT(BaseModel):
+    default_conf = {
+        "rootsift": True,
+        "nms_radius": 0,  # None to disable filtering entirely.
+        "max_keypoints": 4096,
+        "backend": "opencv",  # in {opencv, pycolmap, pycolmap_cpu, pycolmap_cuda}
+        "detection_threshold": 0.0066667,  # from COLMAP
+        "edge_threshold": 10,
+        "first_octave": -1,  # only used by pycolmap, the default of COLMAP
+        "num_octaves": 4,
+    }
+
+    required_data_keys = ["image"]
+
+    def _init(self, conf):
+        self.conf = OmegaConf.create(self.conf)
+        backend = self.conf.backend
+        if backend.startswith("pycolmap"):
+            if pycolmap is None:
+                raise ImportError(
+                    "Cannot find module pycolmap: install it with pip"
+                    "or use backend=opencv."
+                )
+            options = {
+                "peak_threshold": self.conf.detection_threshold,
+                "edge_threshold": self.conf.edge_threshold,
+                "first_octave": self.conf.first_octave,
+                "num_octaves": self.conf.num_octaves,
+                "normalization": pycolmap.Normalization.L2,  # L1_ROOT is buggy.
+            }
+            device = (
+                "auto"
+                if backend == "pycolmap"
+                else backend.replace("pycolmap_", "")
+            )
+            if (
+                backend == "pycolmap_cpu" or not pycolmap.has_cuda
+            ) and pycolmap.__version__ < "0.5.0":
+                warnings.warn(
+                    "The pycolmap CPU SIFT is buggy in version < 0.5.0, "
+                    "consider upgrading pycolmap or use the CUDA version.",
+                    stacklevel=1,
+                )
+            else:
+                options["max_num_features"] = self.conf.max_keypoints
+            self.sift = pycolmap.Sift(options=options, device=device)
+        elif backend == "opencv":
+            self.sift = cv2.SIFT_create(
+                contrastThreshold=self.conf.detection_threshold,
+                nfeatures=self.conf.max_keypoints,
+                edgeThreshold=self.conf.edge_threshold,
+                nOctaveLayers=self.conf.num_octaves,
+            )
+        else:
+            backends = {"opencv", "pycolmap", "pycolmap_cpu", "pycolmap_cuda"}
+            raise ValueError(
+                f"Unknown backend: {backend} not in "
+                f"{{{','.join(backends)}}}."
+            )
+        logger.info("Load SIFT model done.")
+
+    def extract_single_image(self, image: torch.Tensor):
+        image_np = image.cpu().numpy().squeeze(0)
+
+        if self.conf.backend.startswith("pycolmap"):
+            if version.parse(pycolmap.__version__) >= version.parse("0.5.0"):
+                detections, descriptors = self.sift.extract(image_np)
+                scores = None  # Scores are not exposed by COLMAP anymore.
+            else:
+                detections, scores, descriptors = self.sift.extract(image_np)
+            keypoints = detections[:, :2]  # Keep only (x, y).
+            scales, angles = detections[:, -2:].T
+            if scores is not None and (
+                self.conf.backend == "pycolmap_cpu" or not pycolmap.has_cuda
+            ):
+                # Set the scores as a combination of abs. response and scale.
+                scores = np.abs(scores) * scales
+        elif self.conf.backend == "opencv":
+            # TODO: Check if opencv keypoints are already in corner convention
+            keypoints, scores, scales, angles, descriptors = run_opencv_sift(
+                self.sift, (image_np * 255.0).astype(np.uint8)
+            )
+        pred = {
+            "keypoints": keypoints,
+            "scales": scales,
+            "oris": angles,
+            "descriptors": descriptors,
+        }
+        if scores is not None:
+            pred["scores"] = scores
+
+        # sometimes pycolmap returns points outside the image. We remove them
+        if self.conf.backend.startswith("pycolmap"):
+            is_inside = (
+                pred["keypoints"] + 0.5 < np.array([image_np.shape[-2:][::-1]])
+            ).all(-1)
+            pred = {k: v[is_inside] for k, v in pred.items()}
+
+        if self.conf.nms_radius is not None:
+            keep = filter_dog_point(
+                pred["keypoints"],
+                pred["scales"],
+                pred["oris"],
+                image_np.shape,
+                self.conf.nms_radius,
+                scores=pred.get("scores"),
+            )
+            pred = {k: v[keep] for k, v in pred.items()}
+
+        pred = {k: torch.from_numpy(v) for k, v in pred.items()}
+        if scores is not None:
+            # Keep the k keypoints with highest score
+            num_points = self.conf.max_keypoints
+            if num_points is not None and len(pred["keypoints"]) > num_points:
+                indices = torch.topk(pred["scores"], num_points).indices
+                pred = {k: v[indices] for k, v in pred.items()}
+        return pred
+
+    def _forward(self, data: dict) -> dict:
+        image = data["image"]
+        if image.shape[1] == 3:
+            image = rgb_to_grayscale(image)
+        device = image.device
+        image = image.cpu()
+        pred = []
+        for k in range(len(image)):
+            img = image[k]
+            if "image_size" in data.keys():
+                # avoid extracting points in padded areas
+                w, h = data["image_size"][k]
+                img = img[:, :h, :w]
+            p = self.extract_single_image(img)
+            pred.append(p)
+        pred = {
+            k: torch.stack([p[k] for p in pred], 0).to(device) for k in pred[0]
+        }
+        if self.conf.rootsift:
+            pred["descriptors"] = sift_to_rootsift(pred["descriptors"])
+        pred["descriptors"] = pred["descriptors"].permute(0, 2, 1)
+        pred["keypoint_scores"] = pred["scores"].clone()
+        return pred

hloc/extractors/superpoint.py

@@ -0,0 +1,51 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+sys.path.append(str(Path(__file__).parent / "../../third_party"))
+from SuperGluePretrainedNetwork.models import superpoint  # noqa E402
+
+
+# The original keypoint sampling is incorrect. We patch it here but
+# we don't fix it upstream to not impact exisiting evaluations.
+def sample_descriptors_fix_sampling(keypoints, descriptors, s: int = 8):
+    """Interpolate descriptors at keypoint locations"""
+    b, c, h, w = descriptors.shape
+    keypoints = (keypoints + 0.5) / (keypoints.new_tensor([w, h]) * s)
+    keypoints = keypoints * 2 - 1  # normalize to (-1, 1)
+    descriptors = torch.nn.functional.grid_sample(
+        descriptors,
+        keypoints.view(b, 1, -1, 2),
+        mode="bilinear",
+        align_corners=False,
+    )
+    descriptors = torch.nn.functional.normalize(
+        descriptors.reshape(b, c, -1), p=2, dim=1
+    )
+    return descriptors
+
+
+class SuperPoint(BaseModel):
+    default_conf = {
+        "nms_radius": 4,
+        "keypoint_threshold": 0.005,
+        "max_keypoints": -1,
+        "remove_borders": 4,
+        "fix_sampling": False,
+    }
+    required_inputs = ["image"]
+    detection_noise = 2.0
+
+    def _init(self, conf):
+        if conf["fix_sampling"]:
+            superpoint.sample_descriptors = sample_descriptors_fix_sampling
+        self.net = superpoint.SuperPoint(conf)
+        logger.info("Load SuperPoint model done.")
+
+    def _forward(self, data):
+        return self.net(data, self.conf)

hloc/extractors/xfeat.py

@@ -0,0 +1,33 @@
+import torch
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+
+class XFeat(BaseModel):
+    default_conf = {
+        "keypoint_threshold": 0.005,
+        "max_keypoints": -1,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "verlab/accelerated_features",
+            "XFeat",
+            pretrained=True,
+            top_k=self.conf["max_keypoints"],
+        )
+        logger.info("Load XFeat(sparse) model done.")
+
+    def _forward(self, data):
+        pred = self.net.detectAndCompute(
+            data["image"], top_k=self.conf["max_keypoints"]
+        )[0]
+        pred = {
+            "keypoints": pred["keypoints"][None],
+            "scores": pred["scores"][None],
+            "descriptors": pred["descriptors"].T[None],
+        }
+        return pred

hloc/localize_inloc.py

@@ -0,0 +1,183 @@
+import argparse
+import pickle
+from pathlib import Path
+
+import cv2
+import h5py
+import numpy as np
+import pycolmap
+import torch
+from scipy.io import loadmat
+from tqdm import tqdm
+
+from . import logger
+from .utils.parsers import names_to_pair, parse_retrieval
+
+
+def interpolate_scan(scan, kp):
+    h, w, c = scan.shape
+    kp = kp / np.array([[w - 1, h - 1]]) * 2 - 1
+    assert np.all(kp > -1) and np.all(kp < 1)
+    scan = torch.from_numpy(scan).permute(2, 0, 1)[None]
+    kp = torch.from_numpy(kp)[None, None]
+    grid_sample = torch.nn.functional.grid_sample
+
+    # To maximize the number of points that have depth:
+    # do bilinear interpolation first and then nearest for the remaining points
+    interp_lin = grid_sample(scan, kp, align_corners=True, mode="bilinear")[
+        0, :, 0
+    ]
+    interp_nn = torch.nn.functional.grid_sample(
+        scan, kp, align_corners=True, mode="nearest"
+    )[0, :, 0]
+    interp = torch.where(torch.isnan(interp_lin), interp_nn, interp_lin)
+    valid = ~torch.any(torch.isnan(interp), 0)
+
+    kp3d = interp.T.numpy()
+    valid = valid.numpy()
+    return kp3d, valid
+
+
+def get_scan_pose(dataset_dir, rpath):
+    split_image_rpath = rpath.split("/")
+    floor_name = split_image_rpath[-3]
+    scan_id = split_image_rpath[-2]
+    image_name = split_image_rpath[-1]
+    building_name = image_name[:3]
+
+    path = Path(
+        dataset_dir,
+        "database/alignments",
+        floor_name,
+        f"transformations/{building_name}_trans_{scan_id}.txt",
+    )
+    with open(path) as f:
+        raw_lines = f.readlines()
+
+    P_after_GICP = np.array(
+        [
+            np.fromstring(raw_lines[7], sep=" "),
+            np.fromstring(raw_lines[8], sep=" "),
+            np.fromstring(raw_lines[9], sep=" "),
+            np.fromstring(raw_lines[10], sep=" "),
+        ]
+    )
+
+    return P_after_GICP
+
+
+def pose_from_cluster(
+    dataset_dir, q, retrieved, feature_file, match_file, skip=None
+):
+    height, width = cv2.imread(str(dataset_dir / q)).shape[:2]
+    cx = 0.5 * width
+    cy = 0.5 * height
+    focal_length = 4032.0 * 28.0 / 36.0
+
+    all_mkpq = []
+    all_mkpr = []
+    all_mkp3d = []
+    all_indices = []
+    kpq = feature_file[q]["keypoints"].__array__()
+    num_matches = 0
+
+    for i, r in enumerate(retrieved):
+        kpr = feature_file[r]["keypoints"].__array__()
+        pair = names_to_pair(q, r)
+        m = match_file[pair]["matches0"].__array__()
+        v = m > -1
+
+        if skip and (np.count_nonzero(v) < skip):
+            continue
+
+        mkpq, mkpr = kpq[v], kpr[m[v]]
+        num_matches += len(mkpq)
+
+        scan_r = loadmat(Path(dataset_dir, r + ".mat"))["XYZcut"]
+        mkp3d, valid = interpolate_scan(scan_r, mkpr)
+        Tr = get_scan_pose(dataset_dir, r)
+        mkp3d = (Tr[:3, :3] @ mkp3d.T + Tr[:3, -1:]).T
+
+        all_mkpq.append(mkpq[valid])
+        all_mkpr.append(mkpr[valid])
+        all_mkp3d.append(mkp3d[valid])
+        all_indices.append(np.full(np.count_nonzero(valid), i))
+
+    all_mkpq = np.concatenate(all_mkpq, 0)
+    all_mkpr = np.concatenate(all_mkpr, 0)
+    all_mkp3d = np.concatenate(all_mkp3d, 0)
+    all_indices = np.concatenate(all_indices, 0)
+
+    cfg = {
+        "model": "SIMPLE_PINHOLE",
+        "width": width,
+        "height": height,
+        "params": [focal_length, cx, cy],
+    }
+    ret = pycolmap.absolute_pose_estimation(all_mkpq, all_mkp3d, cfg, 48.00)
+    ret["cfg"] = cfg
+    return ret, all_mkpq, all_mkpr, all_mkp3d, all_indices, num_matches
+
+
+def main(dataset_dir, retrieval, features, matches, results, skip_matches=None):
+    assert retrieval.exists(), retrieval
+    assert features.exists(), features
+    assert matches.exists(), matches
+
+    retrieval_dict = parse_retrieval(retrieval)
+    queries = list(retrieval_dict.keys())
+
+    feature_file = h5py.File(features, "r", libver="latest")
+    match_file = h5py.File(matches, "r", libver="latest")
+
+    poses = {}
+    logs = {
+        "features": features,
+        "matches": matches,
+        "retrieval": retrieval,
+        "loc": {},
+    }
+    logger.info("Starting localization...")
+    for q in tqdm(queries):
+        db = retrieval_dict[q]
+        ret, mkpq, mkpr, mkp3d, indices, num_matches = pose_from_cluster(
+            dataset_dir, q, db, feature_file, match_file, skip_matches
+        )
+
+        poses[q] = (ret["qvec"], ret["tvec"])
+        logs["loc"][q] = {
+            "db": db,
+            "PnP_ret": ret,
+            "keypoints_query": mkpq,
+            "keypoints_db": mkpr,
+            "3d_points": mkp3d,
+            "indices_db": indices,
+            "num_matches": num_matches,
+        }
+
+    logger.info(f"Writing poses to {results}...")
+    with open(results, "w") as f:
+        for q in queries:
+            qvec, tvec = poses[q]
+            qvec = " ".join(map(str, qvec))
+            tvec = " ".join(map(str, tvec))
+            name = q.split("/")[-1]
+            f.write(f"{name} {qvec} {tvec}\n")
+
+    logs_path = f"{results}_logs.pkl"
+    logger.info(f"Writing logs to {logs_path}...")
+    with open(logs_path, "wb") as f:
+        pickle.dump(logs, f)
+    logger.info("Done!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_dir", type=Path, required=True)
+    parser.add_argument("--retrieval", type=Path, required=True)
+    parser.add_argument("--features", type=Path, required=True)
+    parser.add_argument("--matches", type=Path, required=True)
+    parser.add_argument("--results", type=Path, required=True)
+    parser.add_argument("--skip_matches", type=int)
+    args = parser.parse_args()
+    main(**args.__dict__)

hloc/localize_sfm.py

@@ -0,0 +1,247 @@
+import argparse
+import pickle
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Union
+
+import numpy as np
+import pycolmap
+from tqdm import tqdm
+
+from . import logger
+from .utils.io import get_keypoints, get_matches
+from .utils.parsers import parse_image_lists, parse_retrieval
+
+
+def do_covisibility_clustering(
+    frame_ids: List[int], reconstruction: pycolmap.Reconstruction
+):
+    clusters = []
+    visited = set()
+    for frame_id in frame_ids:
+        # Check if already labeled
+        if frame_id in visited:
+            continue
+
+        # New component
+        clusters.append([])
+        queue = {frame_id}
+        while len(queue):
+            exploration_frame = queue.pop()
+
+            # Already part of the component
+            if exploration_frame in visited:
+                continue
+            visited.add(exploration_frame)
+            clusters[-1].append(exploration_frame)
+
+            observed = reconstruction.images[exploration_frame].points2D
+            connected_frames = {
+                obs.image_id
+                for p2D in observed
+                if p2D.has_point3D()
+                for obs in reconstruction.points3D[
+                    p2D.point3D_id
+                ].track.elements
+            }
+            connected_frames &= set(frame_ids)
+            connected_frames -= visited
+            queue |= connected_frames
+
+    clusters = sorted(clusters, key=len, reverse=True)
+    return clusters
+
+
+class QueryLocalizer:
+    def __init__(self, reconstruction, config=None):
+        self.reconstruction = reconstruction
+        self.config = config or {}
+
+    def localize(self, points2D_all, points2D_idxs, points3D_id, query_camera):
+        points2D = points2D_all[points2D_idxs]
+        points3D = [self.reconstruction.points3D[j].xyz for j in points3D_id]
+        ret = pycolmap.absolute_pose_estimation(
+            points2D,
+            points3D,
+            query_camera,
+            estimation_options=self.config.get("estimation", {}),
+            refinement_options=self.config.get("refinement", {}),
+        )
+        return ret
+
+
+def pose_from_cluster(
+    localizer: QueryLocalizer,
+    qname: str,
+    query_camera: pycolmap.Camera,
+    db_ids: List[int],
+    features_path: Path,
+    matches_path: Path,
+    **kwargs,
+):
+    kpq = get_keypoints(features_path, qname)
+    kpq += 0.5  # COLMAP coordinates
+
+    kp_idx_to_3D = defaultdict(list)
+    kp_idx_to_3D_to_db = defaultdict(lambda: defaultdict(list))
+    num_matches = 0
+    for i, db_id in enumerate(db_ids):
+        image = localizer.reconstruction.images[db_id]
+        if image.num_points3D == 0:
+            logger.debug(f"No 3D points found for {image.name}.")
+            continue
+        points3D_ids = np.array(
+            [p.point3D_id if p.has_point3D() else -1 for p in image.points2D]
+        )
+
+        matches, _ = get_matches(matches_path, qname, image.name)
+        matches = matches[points3D_ids[matches[:, 1]] != -1]
+        num_matches += len(matches)
+        for idx, m in matches:
+            id_3D = points3D_ids[m]
+            kp_idx_to_3D_to_db[idx][id_3D].append(i)
+            # avoid duplicate observations
+            if id_3D not in kp_idx_to_3D[idx]:
+                kp_idx_to_3D[idx].append(id_3D)
+
+    idxs = list(kp_idx_to_3D.keys())
+    mkp_idxs = [i for i in idxs for _ in kp_idx_to_3D[i]]
+    mp3d_ids = [j for i in idxs for j in kp_idx_to_3D[i]]
+    ret = localizer.localize(kpq, mkp_idxs, mp3d_ids, query_camera, **kwargs)
+    if ret is not None:
+        ret["camera"] = query_camera
+
+    # mostly for logging and post-processing
+    mkp_to_3D_to_db = [
+        (j, kp_idx_to_3D_to_db[i][j]) for i in idxs for j in kp_idx_to_3D[i]
+    ]
+    log = {
+        "db": db_ids,
+        "PnP_ret": ret,
+        "keypoints_query": kpq[mkp_idxs],
+        "points3D_ids": mp3d_ids,
+        "points3D_xyz": None,  # we don't log xyz anymore because of file size
+        "num_matches": num_matches,
+        "keypoint_index_to_db": (mkp_idxs, mkp_to_3D_to_db),
+    }
+    return ret, log
+
+
+def main(
+    reference_sfm: Union[Path, pycolmap.Reconstruction],
+    queries: Path,
+    retrieval: Path,
+    features: Path,
+    matches: Path,
+    results: Path,
+    ransac_thresh: int = 12,
+    covisibility_clustering: bool = False,
+    prepend_camera_name: bool = False,
+    config: Dict = None,
+):
+    assert retrieval.exists(), retrieval
+    assert features.exists(), features
+    assert matches.exists(), matches
+
+    queries = parse_image_lists(queries, with_intrinsics=True)
+    retrieval_dict = parse_retrieval(retrieval)
+
+    logger.info("Reading the 3D model...")
+    if not isinstance(reference_sfm, pycolmap.Reconstruction):
+        reference_sfm = pycolmap.Reconstruction(reference_sfm)
+    db_name_to_id = {img.name: i for i, img in reference_sfm.images.items()}
+
+    config = {
+        "estimation": {"ransac": {"max_error": ransac_thresh}},
+        **(config or {}),
+    }
+    localizer = QueryLocalizer(reference_sfm, config)
+
+    cam_from_world = {}
+    logs = {
+        "features": features,
+        "matches": matches,
+        "retrieval": retrieval,
+        "loc": {},
+    }
+    logger.info("Starting localization...")
+    for qname, qcam in tqdm(queries):
+        if qname not in retrieval_dict:
+            logger.warning(
+                f"No images retrieved for query image {qname}. Skipping..."
+            )
+            continue
+        db_names = retrieval_dict[qname]
+        db_ids = []
+        for n in db_names:
+            if n not in db_name_to_id:
+                logger.warning(f"Image {n} was retrieved but not in database")
+                continue
+            db_ids.append(db_name_to_id[n])
+
+        if covisibility_clustering:
+            clusters = do_covisibility_clustering(db_ids, reference_sfm)
+            best_inliers = 0
+            best_cluster = None
+            logs_clusters = []
+            for i, cluster_ids in enumerate(clusters):
+                ret, log = pose_from_cluster(
+                    localizer, qname, qcam, cluster_ids, features, matches
+                )
+                if ret is not None and ret["num_inliers"] > best_inliers:
+                    best_cluster = i
+                    best_inliers = ret["num_inliers"]
+                logs_clusters.append(log)
+            if best_cluster is not None:
+                ret = logs_clusters[best_cluster]["PnP_ret"]
+                cam_from_world[qname] = ret["cam_from_world"]
+            logs["loc"][qname] = {
+                "db": db_ids,
+                "best_cluster": best_cluster,
+                "log_clusters": logs_clusters,
+                "covisibility_clustering": covisibility_clustering,
+            }
+        else:
+            ret, log = pose_from_cluster(
+                localizer, qname, qcam, db_ids, features, matches
+            )
+            if ret is not None:
+                cam_from_world[qname] = ret["cam_from_world"]
+            else:
+                closest = reference_sfm.images[db_ids[0]]
+                cam_from_world[qname] = closest.cam_from_world
+            log["covisibility_clustering"] = covisibility_clustering
+            logs["loc"][qname] = log
+
+    logger.info(f"Localized {len(cam_from_world)} / {len(queries)} images.")
+    logger.info(f"Writing poses to {results}...")
+    with open(results, "w") as f:
+        for query, t in cam_from_world.items():
+            qvec = " ".join(map(str, t.rotation.quat[[3, 0, 1, 2]]))
+            tvec = " ".join(map(str, t.translation))
+            name = query.split("/")[-1]
+            if prepend_camera_name:
+                name = query.split("/")[-2] + "/" + name
+            f.write(f"{name} {qvec} {tvec}\n")
+
+    logs_path = f"{results}_logs.pkl"
+    logger.info(f"Writing logs to {logs_path}...")
+    # TODO: Resolve pickling issue with pycolmap objects.
+    with open(logs_path, "wb") as f:
+        pickle.dump(logs, f)
+    logger.info("Done!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--reference_sfm", type=Path, required=True)
+    parser.add_argument("--queries", type=Path, required=True)
+    parser.add_argument("--features", type=Path, required=True)
+    parser.add_argument("--matches", type=Path, required=True)
+    parser.add_argument("--retrieval", type=Path, required=True)
+    parser.add_argument("--results", type=Path, required=True)
+    parser.add_argument("--ransac_thresh", type=float, default=12.0)
+    parser.add_argument("--covisibility_clustering", action="store_true")
+    parser.add_argument("--prepend_camera_name", action="store_true")
+    args = parser.parse_args()
+    main(**args.__dict__)

hloc/match_dense.py

@@ -0,0 +1,1121 @@
+import argparse
+import pprint
+from collections import Counter, defaultdict
+from itertools import chain
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
+
+import cv2
+import h5py
+import numpy as np
+import torch
+import torchvision.transforms.functional as F
+from scipy.spatial import KDTree
+from tqdm import tqdm
+
+from . import logger, matchers
+from .extract_features import read_image, resize_image
+from .match_features import find_unique_new_pairs
+from .utils.base_model import dynamic_load
+from .utils.io import list_h5_names
+from .utils.parsers import names_to_pair, parse_retrieval
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+confs = {
+    # Best quality but loads of points. Only use for small scenes
+    "loftr": {
+        "output": "matches-loftr",
+        "model": {
+            "name": "loftr",
+            "weights": "outdoor",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "width": 640,
+            "height": 480,
+            "force_resize": True,
+        },
+        "max_error": 1,  # max error for assigned keypoints (in px)
+        "cell_size": 1,  # size of quantization patch (max 1 kp/patch)
+    },
+    "eloftr": {
+        "output": "matches-eloftr",
+        "model": {
+            "name": "eloftr",
+            "weights": "weights/eloftr_outdoor.ckpt",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 32,
+            "width": 640,
+            "height": 480,
+            "force_resize": True,
+        },
+        "max_error": 1,  # max error for assigned keypoints (in px)
+        "cell_size": 1,  # size of quantization patch (max 1 kp/patch)
+    },
+    # "loftr_quadtree": {
+    #     "output": "matches-loftr-quadtree",
+    #     "model": {
+    #         "name": "quadtree",
+    #         "weights": "outdoor",
+    #         "max_keypoints": 2000,
+    #         "match_threshold": 0.2,
+    #     },
+    #     "preprocessing": {
+    #         "grayscale": True,
+    #         "resize_max": 1024,
+    #         "dfactor": 8,
+    #         "width": 640,
+    #         "height": 480,
+    #         "force_resize": True,
+    #     },
+    #     "max_error": 1,  # max error for assigned keypoints (in px)
+    #     "cell_size": 1,  # size of quantization patch (max 1 kp/patch)
+    # },
+    "cotr": {
+        "output": "matches-cotr",
+        "model": {
+            "name": "cotr",
+            "weights": "out/default",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "width": 640,
+            "height": 480,
+            "force_resize": True,
+        },
+        "max_error": 1,  # max error for assigned keypoints (in px)
+        "cell_size": 1,  # size of quantization patch (max 1 kp/patch)
+    },
+    # Semi-scalable loftr which limits detected keypoints
+    "loftr_aachen": {
+        "output": "matches-loftr_aachen",
+        "model": {
+            "name": "loftr",
+            "weights": "outdoor",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "width": 640,
+            "height": 480,
+            "force_resize": True,
+        },
+        "max_error": 2,  # max error for assigned keypoints (in px)
+        "cell_size": 8,  # size of quantization patch (max 1 kp/patch)
+    },
+    # Use for matching superpoint feats with loftr
+    "loftr_superpoint": {
+        "output": "matches-loftr_aachen",
+        "model": {
+            "name": "loftr",
+            "weights": "outdoor",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "width": 640,
+            "height": 480,
+            "force_resize": True,
+        },
+        "max_error": 4,  # max error for assigned keypoints (in px)
+        "cell_size": 4,  # size of quantization patch (max 1 kp/patch)
+    },
+    # Use topicfm for matching feats
+    "topicfm": {
+        "output": "matches-topicfm",
+        "model": {
+            "name": "topicfm",
+            "weights": "outdoor",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "width": 640,
+            "height": 480,
+        },
+    },
+    # Use aspanformer for matching feats
+    "aspanformer": {
+        "output": "matches-aspanformer",
+        "model": {
+            "name": "aspanformer",
+            "weights": "outdoor",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "duster": {
+        "output": "matches-duster",
+        "model": {
+            "name": "duster",
+            "weights": "vit_large",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 512,
+            "dfactor": 16,
+        },
+    },
+    "mast3r": {
+        "output": "matches-mast3r",
+        "model": {
+            "name": "mast3r",
+            "weights": "vit_large",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 512,
+            "dfactor": 16,
+        },
+    },
+    "xfeat_lightglue": {
+        "output": "matches-xfeat_lightglue",
+        "model": {
+            "name": "xfeat_lightglue",
+            "max_keypoints": 8000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": False,
+            "resize_max": 1024,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "xfeat_dense": {
+        "output": "matches-xfeat_dense",
+        "model": {
+            "name": "xfeat_dense",
+            "max_keypoints": 8000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": False,
+            "resize_max": 1024,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "dkm": {
+        "output": "matches-dkm",
+        "model": {
+            "name": "dkm",
+            "weights": "outdoor",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 80,
+            "height": 60,
+            "dfactor": 8,
+        },
+    },
+    "roma": {
+        "output": "matches-roma",
+        "model": {
+            "name": "roma",
+            "weights": "outdoor",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 320,
+            "height": 240,
+            "dfactor": 8,
+        },
+    },
+    "gim(dkm)": {
+        "output": "matches-gim",
+        "model": {
+            "name": "gim",
+            "weights": "gim_dkm_100h.ckpt",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 320,
+            "height": 240,
+            "dfactor": 8,
+        },
+    },
+    "omniglue": {
+        "output": "matches-omniglue",
+        "model": {
+            "name": "omniglue",
+            "match_threshold": 0.2,
+            "max_keypoints": 2000,
+            "features": "null",
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+            "resize_max": 1024,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "sold2": {
+        "output": "matches-sold2",
+        "model": {
+            "name": "sold2",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+    "gluestick": {
+        "output": "matches-gluestick",
+        "model": {
+            "name": "gluestick",
+            "use_lines": True,
+            "max_keypoints": 1000,
+            "max_lines": 300,
+            "force_num_keypoints": False,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 640,
+            "height": 480,
+            "dfactor": 8,
+        },
+    },
+}
+
+
+def to_cpts(kpts, ps):
+    if ps > 0.0:
+        kpts = np.round(np.round((kpts + 0.5) / ps) * ps - 0.5, 2)
+    return [tuple(cpt) for cpt in kpts]
+
+
+def assign_keypoints(
+    kpts: np.ndarray,
+    other_cpts: Union[List[Tuple], np.ndarray],
+    max_error: float,
+    update: bool = False,
+    ref_bins: Optional[List[Counter]] = None,
+    scores: Optional[np.ndarray] = None,
+    cell_size: Optional[int] = None,
+):
+    if not update:
+        # Without update this is just a NN search
+        if len(other_cpts) == 0 or len(kpts) == 0:
+            return np.full(len(kpts), -1)
+        dist, kpt_ids = KDTree(np.array(other_cpts)).query(kpts)
+        valid = dist <= max_error
+        kpt_ids[~valid] = -1
+        return kpt_ids
+    else:
+        ps = cell_size if cell_size is not None else max_error
+        ps = max(ps, max_error)
+        # With update we quantize and bin (optionally)
+        assert isinstance(other_cpts, list)
+        kpt_ids = []
+        cpts = to_cpts(kpts, ps)
+        bpts = to_cpts(kpts, int(max_error))
+        cp_to_id = {val: i for i, val in enumerate(other_cpts)}
+        for i, (cpt, bpt) in enumerate(zip(cpts, bpts)):
+            try:
+                kid = cp_to_id[cpt]
+            except KeyError:
+                kid = len(cp_to_id)
+                cp_to_id[cpt] = kid
+                other_cpts.append(cpt)
+                if ref_bins is not None:
+                    ref_bins.append(Counter())
+            if ref_bins is not None:
+                score = scores[i] if scores is not None else 1
+                ref_bins[cp_to_id[cpt]][bpt] += score
+            kpt_ids.append(kid)
+        return np.array(kpt_ids)
+
+
+def get_grouped_ids(array):
+    # Group array indices based on its values
+    # all duplicates are grouped as a set
+    idx_sort = np.argsort(array)
+    sorted_array = array[idx_sort]
+    _, ids, _ = np.unique(sorted_array, return_counts=True, return_index=True)
+    res = np.split(idx_sort, ids[1:])
+    return res
+
+
+def get_unique_matches(match_ids, scores):
+    if len(match_ids.shape) == 1:
+        return [0]
+
+    isets1 = get_grouped_ids(match_ids[:, 0])
+    isets2 = get_grouped_ids(match_ids[:, 1])
+    uid1s = [ids[scores[ids].argmax()] for ids in isets1 if len(ids) > 0]
+    uid2s = [ids[scores[ids].argmax()] for ids in isets2 if len(ids) > 0]
+    uids = list(set(uid1s).intersection(uid2s))
+    return match_ids[uids], scores[uids]
+
+
+def matches_to_matches0(matches, scores):
+    if len(matches) == 0:
+        return np.zeros(0, dtype=np.int32), np.zeros(0, dtype=np.float16)
+    n_kps0 = np.max(matches[:, 0]) + 1
+    matches0 = -np.ones((n_kps0,))
+    scores0 = np.zeros((n_kps0,))
+    matches0[matches[:, 0]] = matches[:, 1]
+    scores0[matches[:, 0]] = scores
+    return matches0.astype(np.int32), scores0.astype(np.float16)
+
+
+def kpids_to_matches0(kpt_ids0, kpt_ids1, scores):
+    valid = (kpt_ids0 != -1) & (kpt_ids1 != -1)
+    matches = np.dstack([kpt_ids0[valid], kpt_ids1[valid]])
+    matches = matches.reshape(-1, 2)
+    scores = scores[valid]
+
+    # Remove n-to-1 matches
+    matches, scores = get_unique_matches(matches, scores)
+    return matches_to_matches0(matches, scores)
+
+
+def scale_keypoints(kpts, scale):
+    if np.any(scale != 1.0):
+        kpts *= kpts.new_tensor(scale)
+    return kpts
+
+
+class ImagePairDataset(torch.utils.data.Dataset):
+    default_conf = {
+        "grayscale": True,
+        "resize_max": 1024,
+        "dfactor": 8,
+        "cache_images": False,
+    }
+
+    def __init__(self, image_dir, conf, pairs):
+        self.image_dir = image_dir
+        self.conf = conf = SimpleNamespace(**{**self.default_conf, **conf})
+        self.pairs = pairs
+        if self.conf.cache_images:
+            image_names = set(sum(pairs, ()))  # unique image names in pairs
+            logger.info(
+                f"Loading and caching {len(image_names)} unique images."
+            )
+            self.images = {}
+            self.scales = {}
+            for name in tqdm(image_names):
+                image = read_image(self.image_dir / name, self.conf.grayscale)
+                self.images[name], self.scales[name] = self.preprocess(image)
+
+    def preprocess(self, image: np.ndarray):
+        image = image.astype(np.float32, copy=False)
+        size = image.shape[:2][::-1]
+        scale = np.array([1.0, 1.0])
+
+        if self.conf.resize_max:
+            scale = self.conf.resize_max / max(size)
+            if scale < 1.0:
+                size_new = tuple(int(round(x * scale)) for x in size)
+                image = resize_image(image, size_new, "cv2_area")
+                scale = np.array(size) / np.array(size_new)
+
+        if self.conf.grayscale:
+            assert image.ndim == 2, image.shape
+            image = image[None]
+        else:
+            image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+        image = torch.from_numpy(image / 255.0).float()
+
+        # assure that the size is divisible by dfactor
+        size_new = tuple(
+            map(
+                lambda x: int(x // self.conf.dfactor * self.conf.dfactor),
+                image.shape[-2:],
+            )
+        )
+        image = F.resize(image, size=size_new)
+        scale = np.array(size) / np.array(size_new)[::-1]
+        return image, scale
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def __getitem__(self, idx):
+        name0, name1 = self.pairs[idx]
+        if self.conf.cache_images:
+            image0, scale0 = self.images[name0], self.scales[name0]
+            image1, scale1 = self.images[name1], self.scales[name1]
+        else:
+            image0 = read_image(self.image_dir / name0, self.conf.grayscale)
+            image1 = read_image(self.image_dir / name1, self.conf.grayscale)
+            image0, scale0 = self.preprocess(image0)
+            image1, scale1 = self.preprocess(image1)
+        return image0, image1, scale0, scale1, name0, name1
+
+
+@torch.no_grad()
+def match_dense(
+    conf: Dict,
+    pairs: List[Tuple[str, str]],
+    image_dir: Path,
+    match_path: Path,  # out
+    existing_refs: Optional[List] = [],
+):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    Model = dynamic_load(matchers, conf["model"]["name"])
+    model = Model(conf["model"]).eval().to(device)
+
+    dataset = ImagePairDataset(image_dir, conf["preprocessing"], pairs)
+    loader = torch.utils.data.DataLoader(
+        dataset, num_workers=16, batch_size=1, shuffle=False
+    )
+
+    logger.info("Performing dense matching...")
+    with h5py.File(str(match_path), "a") as fd:
+        for data in tqdm(loader, smoothing=0.1):
+            # load image-pair data
+            image0, image1, scale0, scale1, (name0,), (name1,) = data
+            scale0, scale1 = scale0[0].numpy(), scale1[0].numpy()
+            image0, image1 = image0.to(device), image1.to(device)
+
+            # match semi-dense
+            # for consistency with pairs_from_*: refine kpts of image0
+            if name0 in existing_refs:
+                # special case: flip to enable refinement in query image
+                pred = model({"image0": image1, "image1": image0})
+                pred = {
+                    **pred,
+                    "keypoints0": pred["keypoints1"],
+                    "keypoints1": pred["keypoints0"],
+                }
+            else:
+                # usual case
+                pred = model({"image0": image0, "image1": image1})
+
+            # Rescale keypoints and move to cpu
+            kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
+            kpts0 = scale_keypoints(kpts0 + 0.5, scale0) - 0.5
+            kpts1 = scale_keypoints(kpts1 + 0.5, scale1) - 0.5
+            kpts0 = kpts0.cpu().numpy()
+            kpts1 = kpts1.cpu().numpy()
+            scores = pred["scores"].cpu().numpy()
+
+            # Write matches and matching scores in hloc format
+            pair = names_to_pair(name0, name1)
+            if pair in fd:
+                del fd[pair]
+            grp = fd.create_group(pair)
+
+            # Write dense matching output
+            grp.create_dataset("keypoints0", data=kpts0)
+            grp.create_dataset("keypoints1", data=kpts1)
+            grp.create_dataset("scores", data=scores)
+    del model, loader
+
+
+# default: quantize all!
+def load_keypoints(
+    conf: Dict, feature_paths_refs: List[Path], quantize: Optional[set] = None
+):
+    name2ref = {
+        n: i for i, p in enumerate(feature_paths_refs) for n in list_h5_names(p)
+    }
+
+    existing_refs = set(name2ref.keys())
+    if quantize is None:
+        quantize = existing_refs  # quantize all
+    if len(existing_refs) > 0:
+        logger.info(f"Loading keypoints from {len(existing_refs)} images.")
+
+    # Load query keypoints
+    cpdict = defaultdict(list)
+    bindict = defaultdict(list)
+    for name in existing_refs:
+        with h5py.File(str(feature_paths_refs[name2ref[name]]), "r") as fd:
+            kps = fd[name]["keypoints"].__array__()
+            if name not in quantize:
+                cpdict[name] = kps
+            else:
+                if "scores" in fd[name].keys():
+                    kp_scores = fd[name]["scores"].__array__()
+                else:
+                    # we set the score to 1.0 if not provided
+                    # increase for more weight on reference keypoints for
+                    # stronger anchoring
+                    kp_scores = [1.0 for _ in range(kps.shape[0])]
+                # bin existing keypoints of reference images for association
+                assign_keypoints(
+                    kps,
+                    cpdict[name],
+                    conf["max_error"],
+                    True,
+                    bindict[name],
+                    kp_scores,
+                    conf["cell_size"],
+                )
+    return cpdict, bindict
+
+
+def aggregate_matches(
+    conf: Dict,
+    pairs: List[Tuple[str, str]],
+    match_path: Path,
+    feature_path: Path,
+    required_queries: Optional[Set[str]] = None,
+    max_kps: Optional[int] = None,
+    cpdict: Dict[str, Iterable] = defaultdict(list),
+    bindict: Dict[str, List[Counter]] = defaultdict(list),
+):
+    if required_queries is None:
+        required_queries = set(sum(pairs, ()))
+        # default: do not overwrite existing features in feature_path!
+        required_queries -= set(list_h5_names(feature_path))
+
+    # if an entry in cpdict is provided as np.ndarray we assume it is fixed
+    required_queries -= set(
+        [k for k, v in cpdict.items() if isinstance(v, np.ndarray)]
+    )
+
+    # sort pairs for reduced RAM
+    pairs_per_q = Counter(list(chain(*pairs)))
+    pairs_score = [min(pairs_per_q[i], pairs_per_q[j]) for i, j in pairs]
+    pairs = [p for _, p in sorted(zip(pairs_score, pairs))]
+
+    if len(required_queries) > 0:
+        logger.info(
+            f"Aggregating keypoints for {len(required_queries)} images."
+        )
+    n_kps = 0
+    with h5py.File(str(match_path), "a") as fd:
+        for name0, name1 in tqdm(pairs, smoothing=0.1):
+            pair = names_to_pair(name0, name1)
+            grp = fd[pair]
+            kpts0 = grp["keypoints0"].__array__()
+            kpts1 = grp["keypoints1"].__array__()
+            scores = grp["scores"].__array__()
+
+            # Aggregate local features
+            update0 = name0 in required_queries
+            update1 = name1 in required_queries
+
+            # in localization we do not want to bin the query kp
+            # assumes that the query is name0!
+            if update0 and not update1 and max_kps is None:
+                max_error0 = cell_size0 = 0.0
+            else:
+                max_error0 = conf["max_error"]
+                cell_size0 = conf["cell_size"]
+
+            # Get match ids and extend query keypoints (cpdict)
+            mkp_ids0 = assign_keypoints(
+                kpts0,
+                cpdict[name0],
+                max_error0,
+                update0,
+                bindict[name0],
+                scores,
+                cell_size0,
+            )
+            mkp_ids1 = assign_keypoints(
+                kpts1,
+                cpdict[name1],
+                conf["max_error"],
+                update1,
+                bindict[name1],
+                scores,
+                conf["cell_size"],
+            )
+
+            # Build matches from assignments
+            matches0, scores0 = kpids_to_matches0(mkp_ids0, mkp_ids1, scores)
+
+            assert kpts0.shape[0] == scores.shape[0]
+            grp.create_dataset("matches0", data=matches0)
+            grp.create_dataset("matching_scores0", data=scores0)
+
+            # Convert bins to kps if finished, and store them
+            for name in (name0, name1):
+                pairs_per_q[name] -= 1
+                if pairs_per_q[name] > 0 or name not in required_queries:
+                    continue
+                kp_score = [c.most_common(1)[0][1] for c in bindict[name]]
+                cpdict[name] = [c.most_common(1)[0][0] for c in bindict[name]]
+                cpdict[name] = np.array(cpdict[name], dtype=np.float32)
+
+                # Select top-k query kps by score (reassign matches later)
+                if max_kps:
+                    top_k = min(max_kps, cpdict[name].shape[0])
+                    top_k = np.argsort(kp_score)[::-1][:top_k]
+                    cpdict[name] = cpdict[name][top_k]
+                    kp_score = np.array(kp_score)[top_k]
+
+                # Write query keypoints
+                with h5py.File(feature_path, "a") as kfd:
+                    if name in kfd:
+                        del kfd[name]
+                    kgrp = kfd.create_group(name)
+                    kgrp.create_dataset("keypoints", data=cpdict[name])
+                    kgrp.create_dataset("score", data=kp_score)
+                    n_kps += cpdict[name].shape[0]
+                del bindict[name]
+
+    if len(required_queries) > 0:
+        avg_kp_per_image = round(n_kps / len(required_queries), 1)
+        logger.info(
+            f"Finished assignment, found {avg_kp_per_image} "
+            f"keypoints/image (avg.), total {n_kps}."
+        )
+    return cpdict
+
+
+def assign_matches(
+    pairs: List[Tuple[str, str]],
+    match_path: Path,
+    keypoints: Union[List[Path], Dict[str, np.array]],
+    max_error: float,
+):
+    if isinstance(keypoints, list):
+        keypoints = load_keypoints({}, keypoints, kpts_as_bin=set([]))
+    assert len(set(sum(pairs, ())) - set(keypoints.keys())) == 0
+    with h5py.File(str(match_path), "a") as fd:
+        for name0, name1 in tqdm(pairs):
+            pair = names_to_pair(name0, name1)
+            grp = fd[pair]
+            kpts0 = grp["keypoints0"].__array__()
+            kpts1 = grp["keypoints1"].__array__()
+            scores = grp["scores"].__array__()
+
+            # NN search across cell boundaries
+            mkp_ids0 = assign_keypoints(kpts0, keypoints[name0], max_error)
+            mkp_ids1 = assign_keypoints(kpts1, keypoints[name1], max_error)
+
+            matches0, scores0 = kpids_to_matches0(mkp_ids0, mkp_ids1, scores)
+
+            # overwrite matches0 and matching_scores0
+            del grp["matches0"], grp["matching_scores0"]
+            grp.create_dataset("matches0", data=matches0)
+            grp.create_dataset("matching_scores0", data=scores0)
+
+
+@torch.no_grad()
+def match_and_assign(
+    conf: Dict,
+    pairs_path: Path,
+    image_dir: Path,
+    match_path: Path,  # out
+    feature_path_q: Path,  # out
+    feature_paths_refs: Optional[List[Path]] = [],
+    max_kps: Optional[int] = 8192,
+    overwrite: bool = False,
+) -> Path:
+    for path in feature_paths_refs:
+        if not path.exists():
+            raise FileNotFoundError(f"Reference feature file {path}.")
+    pairs = parse_retrieval(pairs_path)
+    pairs = [(q, r) for q, rs in pairs.items() for r in rs]
+    pairs = find_unique_new_pairs(pairs, None if overwrite else match_path)
+    required_queries = set(sum(pairs, ()))
+
+    name2ref = {
+        n: i for i, p in enumerate(feature_paths_refs) for n in list_h5_names(p)
+    }
+    existing_refs = required_queries.intersection(set(name2ref.keys()))
+
+    # images which require feature extraction
+    required_queries = required_queries - existing_refs
+
+    if feature_path_q.exists():
+        existing_queries = set(list_h5_names(feature_path_q))
+        feature_paths_refs.append(feature_path_q)
+        existing_refs = set.union(existing_refs, existing_queries)
+        if not overwrite:
+            required_queries = required_queries - existing_queries
+
+    if len(pairs) == 0 and len(required_queries) == 0:
+        logger.info("All pairs exist. Skipping dense matching.")
+        return
+
+    # extract semi-dense matches
+    match_dense(conf, pairs, image_dir, match_path, existing_refs=existing_refs)
+
+    logger.info("Assigning matches...")
+
+    # Pre-load existing keypoints
+    cpdict, bindict = load_keypoints(
+        conf, feature_paths_refs, quantize=required_queries
+    )
+
+    # Reassign matches by aggregation
+    cpdict = aggregate_matches(
+        conf,
+        pairs,
+        match_path,
+        feature_path=feature_path_q,
+        required_queries=required_queries,
+        max_kps=max_kps,
+        cpdict=cpdict,
+        bindict=bindict,
+    )
+
+    # Invalidate matches that are far from selected bin by reassignment
+    if max_kps is not None:
+        logger.info(f'Reassign matches with max_error={conf["max_error"]}.')
+        assign_matches(pairs, match_path, cpdict, max_error=conf["max_error"])
+
+
+def scale_lines(lines, scale):
+    if np.any(scale != 1.0):
+        lines *= lines.new_tensor(scale)
+    return lines
+
+
+def match(model, path_0, path_1, conf):
+    default_conf = {
+        "grayscale": True,
+        "resize_max": 1024,
+        "dfactor": 8,
+        "cache_images": False,
+        "force_resize": False,
+        "width": 320,
+        "height": 240,
+    }
+
+    def preprocess(image: np.ndarray):
+        image = image.astype(np.float32, copy=False)
+        size = image.shape[:2][::-1]
+        scale = np.array([1.0, 1.0])
+        if conf.resize_max:
+            scale = conf.resize_max / max(size)
+            if scale < 1.0:
+                size_new = tuple(int(round(x * scale)) for x in size)
+                image = resize_image(image, size_new, "cv2_area")
+                scale = np.array(size) / np.array(size_new)
+        if conf.force_resize:
+            size = image.shape[:2][::-1]
+            image = resize_image(image, (conf.width, conf.height), "cv2_area")
+            size_new = (conf.width, conf.height)
+            scale = np.array(size) / np.array(size_new)
+        if conf.grayscale:
+            assert image.ndim == 2, image.shape
+            image = image[None]
+        else:
+            image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+        image = torch.from_numpy(image / 255.0).float()
+        # assure that the size is divisible by dfactor
+        size_new = tuple(
+            map(
+                lambda x: int(x // conf.dfactor * conf.dfactor),
+                image.shape[-2:],
+            )
+        )
+        image = F.resize(image, size=size_new, antialias=True)
+        scale = np.array(size) / np.array(size_new)[::-1]
+        return image, scale
+
+    conf = SimpleNamespace(**{**default_conf, **conf})
+    image0 = read_image(path_0, conf.grayscale)
+    image1 = read_image(path_1, conf.grayscale)
+    image0, scale0 = preprocess(image0)
+    image1, scale1 = preprocess(image1)
+    image0 = image0.to(device)[None]
+    image1 = image1.to(device)[None]
+    pred = model({"image0": image0, "image1": image1})
+
+    # Rescale keypoints and move to cpu
+    kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
+    kpts0 = scale_keypoints(kpts0 + 0.5, scale0) - 0.5
+    kpts1 = scale_keypoints(kpts1 + 0.5, scale1) - 0.5
+
+    ret = {
+        "image0": image0.squeeze().cpu().numpy(),
+        "image1": image1.squeeze().cpu().numpy(),
+        "keypoints0": kpts0.cpu().numpy(),
+        "keypoints1": kpts1.cpu().numpy(),
+    }
+    if "mconf" in pred.keys():
+        ret["mconf"] = pred["mconf"].cpu().numpy()
+    return ret
+
+
+@torch.no_grad()
+def match_images(model, image_0, image_1, conf, device="cpu"):
+    default_conf = {
+        "grayscale": True,
+        "resize_max": 1024,
+        "dfactor": 8,
+        "cache_images": False,
+        "force_resize": False,
+        "width": 320,
+        "height": 240,
+    }
+
+    def preprocess(image: np.ndarray):
+        image = image.astype(np.float32, copy=False)
+        size = image.shape[:2][::-1]
+        scale = np.array([1.0, 1.0])
+        if conf.resize_max:
+            scale = conf.resize_max / max(size)
+            if scale < 1.0:
+                size_new = tuple(int(round(x * scale)) for x in size)
+                image = resize_image(image, size_new, "cv2_area")
+                scale = np.array(size) / np.array(size_new)
+        if conf.force_resize:
+            size = image.shape[:2][::-1]
+            image = resize_image(image, (conf.width, conf.height), "cv2_area")
+            size_new = (conf.width, conf.height)
+            scale = np.array(size) / np.array(size_new)
+        if conf.grayscale:
+            assert image.ndim == 2, image.shape
+            image = image[None]
+        else:
+            image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+        image = torch.from_numpy(image / 255.0).float()
+
+        # assure that the size is divisible by dfactor
+        size_new = tuple(
+            map(
+                lambda x: int(x // conf.dfactor * conf.dfactor),
+                image.shape[-2:],
+            )
+        )
+        image = F.resize(image, size=size_new)
+        scale = np.array(size) / np.array(size_new)[::-1]
+        return image, scale
+
+    conf = SimpleNamespace(**{**default_conf, **conf})
+
+    if len(image_0.shape) == 3 and conf.grayscale:
+        image0 = cv2.cvtColor(image_0, cv2.COLOR_RGB2GRAY)
+    else:
+        image0 = image_0
+    if len(image_0.shape) == 3 and conf.grayscale:
+        image1 = cv2.cvtColor(image_1, cv2.COLOR_RGB2GRAY)
+    else:
+        image1 = image_1
+
+    # comment following lines, image is always RGB mode
+    # if not conf.grayscale and len(image0.shape) == 3:
+    #     image0 = image0[:, :, ::-1]  # BGR to RGB
+    # if not conf.grayscale and len(image1.shape) == 3:
+    #     image1 = image1[:, :, ::-1]  # BGR to RGB
+
+    image0, scale0 = preprocess(image0)
+    image1, scale1 = preprocess(image1)
+    image0 = image0.to(device)[None]
+    image1 = image1.to(device)[None]
+    pred = model({"image0": image0, "image1": image1})
+
+    s0 = np.array(image_0.shape[:2][::-1]) / np.array(image0.shape[-2:][::-1])
+    s1 = np.array(image_1.shape[:2][::-1]) / np.array(image1.shape[-2:][::-1])
+
+    # Rescale keypoints and move to cpu
+    if "keypoints0" in pred.keys() and "keypoints1" in pred.keys():
+        kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
+        kpts0_origin = scale_keypoints(kpts0 + 0.5, s0) - 0.5
+        kpts1_origin = scale_keypoints(kpts1 + 0.5, s1) - 0.5
+
+        ret = {
+            "image0": image0.squeeze().cpu().numpy(),
+            "image1": image1.squeeze().cpu().numpy(),
+            "image0_orig": image_0,
+            "image1_orig": image_1,
+            "keypoints0": kpts0.cpu().numpy(),
+            "keypoints1": kpts1.cpu().numpy(),
+            "keypoints0_orig": kpts0_origin.cpu().numpy(),
+            "keypoints1_orig": kpts1_origin.cpu().numpy(),
+            "mkeypoints0": kpts0.cpu().numpy(),
+            "mkeypoints1": kpts1.cpu().numpy(),
+            "mkeypoints0_orig": kpts0_origin.cpu().numpy(),
+            "mkeypoints1_orig": kpts1_origin.cpu().numpy(),
+            "original_size0": np.array(image_0.shape[:2][::-1]),
+            "original_size1": np.array(image_1.shape[:2][::-1]),
+            "new_size0": np.array(image0.shape[-2:][::-1]),
+            "new_size1": np.array(image1.shape[-2:][::-1]),
+            "scale0": s0,
+            "scale1": s1,
+        }
+        if "mconf" in pred.keys():
+            ret["mconf"] = pred["mconf"].cpu().numpy()
+        elif "scores" in pred.keys():  # adapting loftr
+            ret["mconf"] = pred["scores"].cpu().numpy()
+        else:
+            ret["mconf"] = np.ones_like(kpts0.cpu().numpy()[:, 0])
+    if "lines0" in pred.keys() and "lines1" in pred.keys():
+        if "keypoints0" in pred.keys() and "keypoints1" in pred.keys():
+            kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
+            kpts0_origin = scale_keypoints(kpts0 + 0.5, s0) - 0.5
+            kpts1_origin = scale_keypoints(kpts1 + 0.5, s1) - 0.5
+            kpts0_origin = kpts0_origin.cpu().numpy()
+            kpts1_origin = kpts1_origin.cpu().numpy()
+        else:
+            kpts0_origin, kpts1_origin = (
+                None,
+                None,
+            )  # np.zeros([0]), np.zeros([0])
+        lines0, lines1 = pred["lines0"], pred["lines1"]
+        lines0_raw, lines1_raw = pred["raw_lines0"], pred["raw_lines1"]
+
+        lines0_raw = torch.from_numpy(lines0_raw.copy())
+        lines1_raw = torch.from_numpy(lines1_raw.copy())
+        lines0_raw = scale_lines(lines0_raw + 0.5, s0) - 0.5
+        lines1_raw = scale_lines(lines1_raw + 0.5, s1) - 0.5
+
+        lines0 = torch.from_numpy(lines0.copy())
+        lines1 = torch.from_numpy(lines1.copy())
+        lines0 = scale_lines(lines0 + 0.5, s0) - 0.5
+        lines1 = scale_lines(lines1 + 0.5, s1) - 0.5
+
+        ret = {
+            "image0_orig": image_0,
+            "image1_orig": image_1,
+            "line0": lines0_raw.cpu().numpy(),
+            "line1": lines1_raw.cpu().numpy(),
+            "line0_orig": lines0.cpu().numpy(),
+            "line1_orig": lines1.cpu().numpy(),
+            "line_keypoints0_orig": kpts0_origin,
+            "line_keypoints1_orig": kpts1_origin,
+        }
+    del pred
+    torch.cuda.empty_cache()
+    return ret
+
+
+@torch.no_grad()
+def main(
+    conf: Dict,
+    pairs: Path,
+    image_dir: Path,
+    export_dir: Optional[Path] = None,
+    matches: Optional[Path] = None,  # out
+    features: Optional[Path] = None,  # out
+    features_ref: Optional[Path] = None,
+    max_kps: Optional[int] = 8192,
+    overwrite: bool = False,
+) -> Path:
+    logger.info(
+        "Extracting semi-dense features with configuration:"
+        f"\n{pprint.pformat(conf)}"
+    )
+
+    if features is None:
+        features = "feats_"
+
+    if isinstance(features, Path):
+        features_q = features
+        if matches is None:
+            raise ValueError(
+                "Either provide both features and matches as Path"
+                " or both as names."
+            )
+    else:
+        if export_dir is None:
+            raise ValueError(
+                "Provide an export_dir if features and matches"
+                f" are not file paths: {features}, {matches}."
+            )
+        features_q = Path(export_dir, f'{features}{conf["output"]}.h5')
+        if matches is None:
+            matches = Path(export_dir, f'{conf["output"]}_{pairs.stem}.h5')
+
+    if features_ref is None:
+        features_ref = []
+    elif isinstance(features_ref, list):
+        features_ref = list(features_ref)
+    elif isinstance(features_ref, Path):
+        features_ref = [features_ref]
+    else:
+        raise TypeError(str(features_ref))
+
+    match_and_assign(
+        conf,
+        pairs,
+        image_dir,
+        matches,
+        features_q,
+        features_ref,
+        max_kps,
+        overwrite,
+    )
+
+    return features_q, matches
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pairs", type=Path, required=True)
+    parser.add_argument("--image_dir", type=Path, required=True)
+    parser.add_argument("--export_dir", type=Path, required=True)
+    parser.add_argument(
+        "--matches", type=Path, default=confs["loftr"]["output"]
+    )
+    parser.add_argument(
+        "--features", type=str, default="feats_" + confs["loftr"]["output"]
+    )
+    parser.add_argument(
+        "--conf", type=str, default="loftr", choices=list(confs.keys())
+    )
+    args = parser.parse_args()
+    main(
+        confs[args.conf],
+        args.pairs,
+        args.image_dir,
+        args.export_dir,
+        args.matches,
+        args.features,
+    )

hloc/match_features.py

@@ -0,0 +1,441 @@
+import argparse
+import pprint
+from functools import partial
+from pathlib import Path
+from queue import Queue
+from threading import Thread
+from typing import Dict, List, Optional, Tuple, Union
+
+import h5py
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from . import logger, matchers
+from .utils.base_model import dynamic_load
+from .utils.parsers import names_to_pair, names_to_pair_old, parse_retrieval
+
+"""
+A set of standard configurations that can be directly selected from the command
+line using their name. Each is a dictionary with the following entries:
+    - output: the name of the match file that will be generated.
+    - model: the model configuration, as passed to a feature matcher.
+"""
+confs = {
+    "superglue": {
+        "output": "matches-superglue",
+        "model": {
+            "name": "superglue",
+            "weights": "outdoor",
+            "sinkhorn_iterations": 50,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "superglue-fast": {
+        "output": "matches-superglue-it5",
+        "model": {
+            "name": "superglue",
+            "weights": "outdoor",
+            "sinkhorn_iterations": 5,
+            "match_threshold": 0.2,
+        },
+    },
+    "superpoint-lightglue": {
+        "output": "matches-lightglue",
+        "model": {
+            "name": "lightglue",
+            "match_threshold": 0.2,
+            "width_confidence": 0.99,  # for point pruning
+            "depth_confidence": 0.95,  # for early stopping,
+            "features": "superpoint",
+            "model_name": "superpoint_lightglue.pth",
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "disk-lightglue": {
+        "output": "matches-disk-lightglue",
+        "model": {
+            "name": "lightglue",
+            "match_threshold": 0.2,
+            "width_confidence": 0.99,  # for point pruning
+            "depth_confidence": 0.95,  # for early stopping,
+            "features": "disk",
+            "model_name": "disk_lightglue.pth",
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "sift-lightglue": {
+        "output": "matches-sift-lightglue",
+        "model": {
+            "name": "lightglue",
+            "match_threshold": 0.2,
+            "width_confidence": 0.99,  # for point pruning
+            "depth_confidence": 0.95,  # for early stopping,
+            "features": "sift",
+            "add_scale_ori": True,
+            "model_name": "sift_lightglue.pth",
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "sgmnet": {
+        "output": "matches-sgmnet",
+        "model": {
+            "name": "sgmnet",
+            "seed_top_k": [256, 256],
+            "seed_radius_coe": 0.01,
+            "net_channels": 128,
+            "layer_num": 9,
+            "head": 4,
+            "seedlayer": [0, 6],
+            "use_mc_seeding": True,
+            "use_score_encoding": False,
+            "conf_bar": [1.11, 0.1],
+            "sink_iter": [10, 100],
+            "detach_iter": 1000000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "NN-superpoint": {
+        "output": "matches-NN-mutual-dist.7",
+        "model": {
+            "name": "nearest_neighbor",
+            "do_mutual_check": True,
+            "distance_threshold": 0.7,
+            "match_threshold": 0.2,
+        },
+    },
+    "NN-ratio": {
+        "output": "matches-NN-mutual-ratio.8",
+        "model": {
+            "name": "nearest_neighbor",
+            "do_mutual_check": True,
+            "ratio_threshold": 0.8,
+            "match_threshold": 0.2,
+        },
+    },
+    "NN-mutual": {
+        "output": "matches-NN-mutual",
+        "model": {
+            "name": "nearest_neighbor",
+            "do_mutual_check": True,
+            "match_threshold": 0.2,
+        },
+    },
+    "Dual-Softmax": {
+        "output": "matches-Dual-Softmax",
+        "model": {
+            "name": "dual_softmax",
+            "match_threshold": 0.01,
+            "inv_temperature": 20,
+        },
+    },
+    "adalam": {
+        "output": "matches-adalam",
+        "model": {
+            "name": "adalam",
+            "match_threshold": 0.2,
+        },
+    },
+    "imp": {
+        "output": "matches-imp",
+        "model": {
+            "name": "imp",
+            "match_threshold": 0.2,
+        },
+    },
+}
+
+
+class WorkQueue:
+    def __init__(self, work_fn, num_threads=1):
+        self.queue = Queue(num_threads)
+        self.threads = [
+            Thread(target=self.thread_fn, args=(work_fn,))
+            for _ in range(num_threads)
+        ]
+        for thread in self.threads:
+            thread.start()
+
+    def join(self):
+        for thread in self.threads:
+            self.queue.put(None)
+        for thread in self.threads:
+            thread.join()
+
+    def thread_fn(self, work_fn):
+        item = self.queue.get()
+        while item is not None:
+            work_fn(item)
+            item = self.queue.get()
+
+    def put(self, data):
+        self.queue.put(data)
+
+
+class FeaturePairsDataset(torch.utils.data.Dataset):
+    def __init__(self, pairs, feature_path_q, feature_path_r):
+        self.pairs = pairs
+        self.feature_path_q = feature_path_q
+        self.feature_path_r = feature_path_r
+
+    def __getitem__(self, idx):
+        name0, name1 = self.pairs[idx]
+        data = {}
+        with h5py.File(self.feature_path_q, "r") as fd:
+            grp = fd[name0]
+            for k, v in grp.items():
+                data[k + "0"] = torch.from_numpy(v.__array__()).float()
+            # some matchers might expect an image but only use its size
+            data["image0"] = torch.empty((1,) + tuple(grp["image_size"])[::-1])
+        with h5py.File(self.feature_path_r, "r") as fd:
+            grp = fd[name1]
+            for k, v in grp.items():
+                data[k + "1"] = torch.from_numpy(v.__array__()).float()
+            data["image1"] = torch.empty((1,) + tuple(grp["image_size"])[::-1])
+        return data
+
+    def __len__(self):
+        return len(self.pairs)
+
+
+def writer_fn(inp, match_path):
+    pair, pred = inp
+    with h5py.File(str(match_path), "a", libver="latest") as fd:
+        if pair in fd:
+            del fd[pair]
+        grp = fd.create_group(pair)
+        matches = pred["matches0"][0].cpu().short().numpy()
+        grp.create_dataset("matches0", data=matches)
+        if "matching_scores0" in pred:
+            scores = pred["matching_scores0"][0].cpu().half().numpy()
+            grp.create_dataset("matching_scores0", data=scores)
+
+
+def main(
+    conf: Dict,
+    pairs: Path,
+    features: Union[Path, str],
+    export_dir: Optional[Path] = None,
+    matches: Optional[Path] = None,
+    features_ref: Optional[Path] = None,
+    overwrite: bool = False,
+) -> Path:
+    if isinstance(features, Path) or Path(features).exists():
+        features_q = features
+        if matches is None:
+            raise ValueError(
+                "Either provide both features and matches as Path"
+                " or both as names."
+            )
+    else:
+        if export_dir is None:
+            raise ValueError(
+                "Provide an export_dir if features is not"
+                f" a file path: {features}."
+            )
+        features_q = Path(export_dir, features + ".h5")
+        if matches is None:
+            matches = Path(
+                export_dir, f'{features}_{conf["output"]}_{pairs.stem}.h5'
+            )
+
+    if features_ref is None:
+        features_ref = features_q
+    match_from_paths(conf, pairs, matches, features_q, features_ref, overwrite)
+
+    return matches
+
+
+def find_unique_new_pairs(pairs_all: List[Tuple[str]], match_path: Path = None):
+    """Avoid to recompute duplicates to save time."""
+    pairs = set()
+    for i, j in pairs_all:
+        if (j, i) not in pairs:
+            pairs.add((i, j))
+    pairs = list(pairs)
+    if match_path is not None and match_path.exists():
+        with h5py.File(str(match_path), "r", libver="latest") as fd:
+            pairs_filtered = []
+            for i, j in pairs:
+                if (
+                    names_to_pair(i, j) in fd
+                    or names_to_pair(j, i) in fd
+                    or names_to_pair_old(i, j) in fd
+                    or names_to_pair_old(j, i) in fd
+                ):
+                    continue
+                pairs_filtered.append((i, j))
+        return pairs_filtered
+    return pairs
+
+
+@torch.no_grad()
+def match_from_paths(
+    conf: Dict,
+    pairs_path: Path,
+    match_path: Path,
+    feature_path_q: Path,
+    feature_path_ref: Path,
+    overwrite: bool = False,
+) -> Path:
+    logger.info(
+        "Matching local features with configuration:"
+        f"\n{pprint.pformat(conf)}"
+    )
+
+    if not feature_path_q.exists():
+        raise FileNotFoundError(f"Query feature file {feature_path_q}.")
+    if not feature_path_ref.exists():
+        raise FileNotFoundError(f"Reference feature file {feature_path_ref}.")
+    match_path.parent.mkdir(exist_ok=True, parents=True)
+
+    assert pairs_path.exists(), pairs_path
+    pairs = parse_retrieval(pairs_path)
+    pairs = [(q, r) for q, rs in pairs.items() for r in rs]
+    pairs = find_unique_new_pairs(pairs, None if overwrite else match_path)
+    if len(pairs) == 0:
+        logger.info("Skipping the matching.")
+        return
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    Model = dynamic_load(matchers, conf["model"]["name"])
+    model = Model(conf["model"]).eval().to(device)
+
+    dataset = FeaturePairsDataset(pairs, feature_path_q, feature_path_ref)
+    loader = torch.utils.data.DataLoader(
+        dataset, num_workers=5, batch_size=1, shuffle=False, pin_memory=True
+    )
+    writer_queue = WorkQueue(partial(writer_fn, match_path=match_path), 5)
+
+    for idx, data in enumerate(tqdm(loader, smoothing=0.1)):
+        data = {
+            k: v if k.startswith("image") else v.to(device, non_blocking=True)
+            for k, v in data.items()
+        }
+        pred = model(data)
+        pair = names_to_pair(*pairs[idx])
+        writer_queue.put((pair, pred))
+    writer_queue.join()
+    logger.info("Finished exporting matches.")
+
+
+def scale_keypoints(kpts, scale):
+    if np.any(scale != 1.0):
+        kpts *= kpts.new_tensor(scale)
+    return kpts
+
+
+@torch.no_grad()
+def match_images(model, feat0, feat1):
+    # forward pass to match keypoints
+    desc0 = feat0["descriptors"][0]
+    desc1 = feat1["descriptors"][0]
+    if len(desc0.shape) == 2:
+        desc0 = desc0.unsqueeze(0)
+    if len(desc1.shape) == 2:
+        desc1 = desc1.unsqueeze(0)
+    if isinstance(feat0["keypoints"], list):
+        feat0["keypoints"] = feat0["keypoints"][0][None]
+    if isinstance(feat1["keypoints"], list):
+        feat1["keypoints"] = feat1["keypoints"][0][None]
+    input_dict = {
+        "image0": feat0["image"],
+        "keypoints0": feat0["keypoints"],
+        "scores0": feat0["scores"][0].unsqueeze(0),
+        "descriptors0": desc0,
+        "image1": feat1["image"],
+        "keypoints1": feat1["keypoints"],
+        "scores1": feat1["scores"][0].unsqueeze(0),
+        "descriptors1": desc1,
+    }
+    if "scales" in feat0:
+        input_dict = {**input_dict, "scales0": feat0["scales"]}
+    if "scales" in feat1:
+        input_dict = {**input_dict, "scales1": feat1["scales"]}
+    if "oris" in feat0:
+        input_dict = {**input_dict, "oris0": feat0["oris"]}
+    if "oris" in feat1:
+        input_dict = {**input_dict, "oris1": feat1["oris"]}
+    pred = model(input_dict)
+    pred = {
+        k: v.cpu().detach()[0] if isinstance(v, torch.Tensor) else v
+        for k, v in pred.items()
+    }
+    kpts0, kpts1 = (
+        feat0["keypoints"][0].cpu().numpy(),
+        feat1["keypoints"][0].cpu().numpy(),
+    )
+    matches, confid = pred["matches0"], pred["matching_scores0"]
+    # Keep the matching keypoints.
+    valid = matches > -1
+    mkpts0 = kpts0[valid]
+    mkpts1 = kpts1[matches[valid]]
+    mconfid = confid[valid]
+    # rescale the keypoints to their original size
+    s0 = feat0["original_size"] / feat0["size"]
+    s1 = feat1["original_size"] / feat1["size"]
+    kpts0_origin = scale_keypoints(torch.from_numpy(kpts0 + 0.5), s0) - 0.5
+    kpts1_origin = scale_keypoints(torch.from_numpy(kpts1 + 0.5), s1) - 0.5
+
+    mkpts0_origin = scale_keypoints(torch.from_numpy(mkpts0 + 0.5), s0) - 0.5
+    mkpts1_origin = scale_keypoints(torch.from_numpy(mkpts1 + 0.5), s1) - 0.5
+
+    ret = {
+        "image0_orig": feat0["image_orig"],
+        "image1_orig": feat1["image_orig"],
+        "keypoints0": kpts0,
+        "keypoints1": kpts1,
+        "keypoints0_orig": kpts0_origin.numpy(),
+        "keypoints1_orig": kpts1_origin.numpy(),
+        "mkeypoints0": mkpts0,
+        "mkeypoints1": mkpts1,
+        "mkeypoints0_orig": mkpts0_origin.numpy(),
+        "mkeypoints1_orig": mkpts1_origin.numpy(),
+        "mconf": mconfid.numpy(),
+    }
+    del feat0, feat1, desc0, desc1, kpts0, kpts1, kpts0_origin, kpts1_origin
+    torch.cuda.empty_cache()
+
+    return ret
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pairs", type=Path, required=True)
+    parser.add_argument("--export_dir", type=Path)
+    parser.add_argument(
+        "--features", type=str, default="feats-superpoint-n4096-r1024"
+    )
+    parser.add_argument("--matches", type=Path)
+    parser.add_argument(
+        "--conf", type=str, default="superglue", choices=list(confs.keys())
+    )
+    args = parser.parse_args()
+    main(confs[args.conf], args.pairs, args.features, args.export_dir)

hloc/matchers/__init__.py

@@ -0,0 +1,3 @@
+def get_matcher(matcher):
+    mod = __import__(f"{__name__}.{matcher}", fromlist=[""])
+    return getattr(mod, "Model")