Init

CODE_OF_CONDUCT.md

@@ -0,0 +1,128 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+[email protected].
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

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.

MANIFEST.in

@@ -0,0 +1,12 @@
+# logo
+include imcui/assets/logo.webp
+
+recursive-include imcui/ui *.yaml
+recursive-include imcui/api *.yaml
+recursive-include imcui/third_party *.yaml *.cfg *.yml
+
+# ui examples
+recursive-include imcui/datasets *.JPG *.jpg *.png
+
+# model
+recursive-include imcui/third_party/SuperGluePretrainedNetwork *.pth

README.md

@@ -1,12 +1,15 @@
 ---
-title: MatchAnything Det
-emoji: 📊
+title: MatchAnything-Det
+emoji: 🏢
 colorFrom: red
-colorTo: green
+colorTo: blue
 sdk: gradio
-sdk_version: 5.49.1
+python_version: 3.10.13
+sdk_version: 4.44.0
+# sdk_version: 5.23.0
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,27 @@
+# import spaces
+import argparse
+from pathlib import Path
+from imcui.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 / "config/config.yaml",
+        help="config file",
+    )
+    args = parser.parse_args()
+    ImageMatchingApp(args.server_name, args.server_port, config=args.config).run()

config/config.yaml

@@ -0,0 +1,33 @@
+server:
+  name: "0.0.0.0"
+  port: 16010
+
+defaults:
+  setting_threshold: 0.1
+  max_keypoints: 2000
+  keypoint_threshold: 0.05
+  enable_ransac: true
+  ransac_method: CV2_USAC_MAGSAC
+  ransac_reproj_threshold: 8
+  ransac_confidence: 0.999
+  ransac_max_iter: 10000
+  ransac_num_samples: 4
+  match_threshold: 0.2
+  setting_geometry: Homography
+
+matcher_zoo:
+  matchanything:
+    matcher: matchanything
+    dense: true
+    info:
+      name: MatchAnything #dispaly name
+      source: "ZJU3DV"
+      display: true
+
+retrieval_zoo:
+  netvlad:
+    enable: true
+  openibl:
+    enable: true
+  cosplace:
+    enable: true

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

environment.yaml

@@ -0,0 +1,14 @@
+name: imw
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.8
+  - pytorch-cuda=11.7
+  - pytorch=1.12.0
+  - torchvision=0.13.1
+  - pip
+  - pip:
+      - -r requirements.txt

imcui/api/__init__.py

@@ -0,0 +1,47 @@
+import base64
+import io
+from typing import List
+
+import numpy as np
+from fastapi.exceptions import HTTPException
+from PIL import Image
+from pydantic import BaseModel
+
+from ..hloc import logger
+from .core import ImageMatchingAPI
+
+
+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
+
+
+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")
+
+
+__all__ = [
+    "ImageMatchingAPI",
+    "ImagesInput",
+    "decode_base64_to_image",
+    "to_base64_nparray",
+]

imcui/api/client.py

@@ -0,0 +1,232 @@
+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,
+    )
+    # breakpoint()
+    # 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__":
+    from pathlib import Path
+
+    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=str(
+            Path(__file__).parents[1]
+            / "datasets/sacre_coeur/mapping_rot/02928139_3448003521_rot45.jpg"
+        ),
+    )
+    parser.add_argument(
+        "--image1",
+        required=False,
+        help="Path for the file's melody",
+        default=str(
+            Path(__file__).parents[1]
+            / "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(1000):
+        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)

imcui/api/config/api.yaml

@@ -0,0 +1,51 @@
+# This file was generated using the `serve build` command on Ray v2.38.0.
+
+proxy_location: EveryNode
+http_options:
+  host: 0.0.0.0
+  port: 8001
+
+grpc_options:
+  port: 9000
+  grpc_servicer_functions: []
+
+logging_config:
+  encoding: TEXT
+  log_level: INFO
+  logs_dir: null
+  enable_access_log: true
+
+applications:
+- name: app1
+  route_prefix: /
+  import_path: api.server:service
+  runtime_env: {}
+  deployments:
+  - name: ImageMatchingService
+    num_replicas: 4
+    ray_actor_options:
+      num_cpus: 2.0
+      num_gpus: 1.0
+
+api:
+  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

imcui/api/core.py

@@ -0,0 +1,308 @@
+# api.py
+import warnings
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from ..hloc import extract_features, logger, match_dense, match_features
+from ..hloc.utils.viz import add_text, plot_keypoints
+from ..ui.utils import filter_matches, get_feature_model, get_model
+from ..ui.viz import display_matches, fig2im, plot_images
+
+warnings.simplefilter("ignore")
+
+
+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
+
+    def _convert_pred(self, pred):
+        ret = {
+            k: v.cpu().detach()[0].numpy() if isinstance(v, torch.Tensor) else v
+            for k, v in pred.items()
+        }
+        ret = {
+            k: v[0].cpu().detach().numpy() if isinstance(v, list) else v
+            for k, v in ret.items()
+        }
+        return ret
+
+    @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 = self._convert_pred(pred)
+        # 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")

imcui/api/server.py

@@ -0,0 +1,170 @@
+# server.py
+import warnings
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import ray
+import torch
+import yaml
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+from PIL import Image
+from ray import serve
+
+from . import ImagesInput, to_base64_nparray
+from .core import ImageMatchingAPI
+from ..hloc import DEVICE
+from ..ui import get_version
+
+warnings.simplefilter("ignore")
+app = FastAPI()
+if ray.is_initialized():
+    ray.shutdown()
+ray.init(
+    dashboard_port=8265,
+    ignore_reinit_error=True,
+)
+serve.start(
+    http_options={"host": "0.0.0.0", "port": 8001},
+)
+
+num_gpus = 1 if torch.cuda.is_available() else 0
+
+
+@serve.deployment(
+    num_replicas=4, ray_actor_options={"num_cpus": 2, "num_gpus": num_gpus}
+)
+@serve.ingress(app)
+class ImageMatchingService:
+    def __init__(self, conf: dict, device: str):
+        self.conf = conf
+        self.api = ImageMatchingAPI(conf=conf, device=device)
+
+    @app.get("/")
+    def root(self):
+        return "Hello, world!"
+
+    @app.get("/version")
+    async def version(self):
+        return {"version": get_version()}
+
+    @app.post("/v1/match")
+    async def match(
+        self, 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)
+
+    @app.post("/v1/extract")
+    async def extract(self, 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):
+        import uvicorn
+
+        uvicorn.run(app, host=host, port=port)
+
+
+def read_config(config_path: Path) -> dict:
+    with open(config_path, "r") as f:
+        conf = yaml.safe_load(f)
+    return conf
+
+
+# api server
+conf = read_config(Path(__file__).parent / "config/api.yaml")
+service = ImageMatchingService.bind(conf=conf["api"], device=DEVICE)
+handle = serve.run(service, route_prefix="/")
+
+# serve run api.server_ray:service
+
+# build to generate config file
+# serve build api.server_ray:service -o api/config/ray.yaml
+# serve run api/config/ray.yaml

imcui/api/test/CMakeLists.txt

@@ -0,0 +1,17 @@
+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})

imcui/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 ..

imcui/api/test/client.cpp

@@ -0,0 +1,81 @@
+#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;
+}

imcui/api/test/helper.h

@@ -0,0 +1,405 @@
+
+#include <b64/encode.h>
+#include <fstream>
+#include <jsoncpp/json/json.h>
+#include <opencv2/opencv.hpp>
+#include <sstream>
+#include <vector>
+
+// base64 to image
+#include <boost/archive/iterators/base64_from_binary.hpp>
+#include <boost/archive/iterators/binary_from_base64.hpp>
+#include <boost/archive/iterators/transform_width.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;
+}

imcui/hloc/__init__.py

@@ -0,0 +1,65 @@
+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")
+
+MODEL_REPO_ID = "LittleFrog/MatchAnything_checkpoints"

imcui/hloc/colmap_from_nvm.py

@@ -0,0 +1,216 @@
+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__)

imcui/hloc/extract_features.py

@@ -0,0 +1,607 @@
+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,
+        },
+    },
+    "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,
+        },
+    },
+    "d2net-ms": {
+        "output": "feats-d2net-ms-n5000-r1600",
+        "model": {
+            "name": "d2net",
+            "multiscale": True,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+        },
+    },
+    "rord": {
+        "output": "feats-rord-ss-n5000-r1600",
+        "model": {
+            "name": "rord",
+            "multiscale": False,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+        },
+    },
+    "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,
+        },
+    },
+    "xfeat": {
+        "output": "feats-xfeat-n5000-r1600",
+        "model": {
+            "name": "xfeat",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+        },
+    },
+    "aliked-n16-rot": {
+        "output": "feats-aliked-n16-rot",
+        "model": {
+            "name": "aliked",
+            "model_name": "aliked-n16rot",
+            "max_num_keypoints": -1,
+            "detection_threshold": 0.2,
+            "nms_radius": 2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1024,
+        },
+    },
+    "aliked-n16": {
+        "output": "feats-aliked-n16",
+        "model": {
+            "name": "aliked",
+            "model_name": "aliked-n16",
+            "max_num_keypoints": -1,
+            "detection_threshold": 0.2,
+            "nms_radius": 2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1024,
+        },
+    },
+    "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,
+        },
+    },
+    "lanet": {
+        "output": "feats-lanet-n5000-r1600",
+        "model": {
+            "name": "lanet",
+            "keypoint_threshold": 0.1,
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+        },
+    },
+    "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)

imcui/hloc/extractors/alike.py

@@ -0,0 +1,61 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, 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):
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}.pth".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+        logger.info("Loaded Alike model from {}".format(model_path))
+        configs[conf["model_name"]]["model_path"] = model_path
+        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],
+        }

imcui/hloc/extractors/aliked.py

@@ -0,0 +1,32 @@
+import sys
+from pathlib import Path
+
+from ..utils.base_model import BaseModel
+
+lightglue_path = Path(__file__).parent / "../../third_party/LightGlue"
+sys.path.append(str(lightglue_path))
+
+from lightglue import ALIKED as ALIKED_
+
+
+class ALIKED(BaseModel):
+    default_conf = {
+        "model_name": "aliked-n16",
+        "max_num_keypoints": -1,
+        "detection_threshold": 0.2,
+        "nms_radius": 2,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        conf.pop("name")
+        self.model = ALIKED_(**conf)
+
+    def _forward(self, data):
+        features = self.model(data)
+
+        return {
+            "keypoints": [f for f in features["keypoints"]],
+            "scores": [f for f in features["keypoint_scores"]],
+            "descriptors": [f.t() for f in features["descriptors"]],
+        }

imcui/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,
+        }

imcui/hloc/extractors/d2net.py

@@ -0,0 +1,60 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, 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):
+        logger.info("Loading D2Net model...")
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+        logger.info(f"Loading model from {model_path}...")
+        self.net = _D2Net(
+            model_file=model_path, 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],
+        }

imcui/hloc/extractors/darkfeat.py

@@ -0,0 +1,44 @@
+import sys
+from pathlib import Path
+
+from .. import MODEL_REPO_ID, 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,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+        logger.info("Loaded DarkFeat model: {}".format(model_path))
+        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
+        }

imcui/hloc/extractors/dedode.py

@@ -0,0 +1,86 @@
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.transforms as transforms
+
+from .. import MODEL_REPO_ID, 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",
+    ]
+
+    # Initialize the line matcher
+    def _init(self, conf):
+        model_detector_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, conf["model_detector_name"]),
+        )
+        model_descriptor_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, conf["model_descriptor_name"]),
+        )
+        logger.info("Loaded DarkFeat model: {}".format(model_detector_path))
+        self.normalizer = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+        # 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
+        }

imcui/hloc/extractors/dir.py

@@ -0,0 +1,78 @@
+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):
+        # todo: download from google drive -> huggingface models
+        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,
+        }

imcui/hloc/extractors/disk.py

@@ -0,0 +1,35 @@
+import kornia
+
+from .. 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],
+        }

imcui/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],
+        }

imcui/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,
+        }

imcui/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,
+        }

imcui/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}

imcui/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
+        }

imcui/hloc/extractors/lanet.py

@@ -0,0 +1,63 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, 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": "PointModel_v0.pth",
+        "keypoint_threshold": 0.1,
+        "max_keypoints": 1024,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        logger.info("Loading LANet model...")
+
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+        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],
+        }

imcui/hloc/extractors/netvlad.py

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

imcui/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,
+        }

imcui/hloc/extractors/r2d2.py

@@ -0,0 +1,73 @@
+import sys
+from pathlib import Path
+
+import torchvision.transforms as tvf
+
+from .. import MODEL_REPO_ID, logger
+
+from ..utils.base_model import BaseModel
+
+r2d2_path = Path(__file__).parents[2] / "third_party/r2d2"
+sys.path.append(str(r2d2_path))
+
+gim_path = Path(__file__).parents[2] / "third_party/gim"
+if str(gim_path) in sys.path:
+    sys.path.remove(str(gim_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_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.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_path)
+        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

imcui/hloc/extractors/rekd.py

@@ -0,0 +1,60 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, 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):
+        # TODO: download model
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+        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],
+        }