add: matching API

README.md

@@ -38,7 +38,8 @@ The tool currently supports various popular image matching algorithms, namely:
 - [x] [RoMa](https://github.com/Vincentqyw/RoMa), CVPR 2024
 - [x] [DeDoDe](https://github.com/Parskatt/DeDoDe), 3DV 2024
 - [ ] [Mickey](https://github.com/nianticlabs/mickey), CVPR 2024
-- [ ] [GIM](https://github.com/xuelunshen/gim), ICLR 2024
+- [x] [GIM](https://github.com/xuelunshen/gim), ICLR 2024
+- [ ] [DUSt3R](https://github.com/naver/dust3r), arXiv 2023
 - [x] [LightGlue](https://github.com/cvg/LightGlue), ICCV 2023
 - [x] [DarkFeat](https://github.com/THU-LYJ-Lab/DarkFeat), AAAI 2023
 - [ ] [ASTR](https://github.com/ASTR2023/ASTR), CVPR 2023

common/api.py

@@ -0,0 +1,298 @@
+import cv2
+import torch
+import warnings
+import numpy as np
+from pathlib import Path
+from typing import Dict, Any, Optional, Tuple, List, Union
+from hloc import logger
+from hloc import match_dense, match_features, extract_features
+from hloc.utils.viz import add_text, plot_keypoints
+from .utils import (
+    load_config,
+    get_model,
+    get_feature_model,
+    filter_matches,
+    device,
+    ROOT,
+)
+from .viz import (
+    fig2im,
+    plot_images,
+    display_matches,
+)
+import matplotlib.pyplot as plt
+
+warnings.simplefilter("ignore")
+
+
+class ImageMatchingAPI(torch.nn.Module):
+    default_conf = {
+        "dense": True,
+        "matcher": {
+            "model": {
+                "name": "topicfm",
+                "match_threshold": 0.2,
+            }
+        },
+        "feature": {
+            "model": {
+                "name": "xfeat",
+                "max_keypoints": 1024,
+                "keypoint_threshold": 0.015,
+            }
+        },
+        "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 = conf = {
+            **self.parse_match_config(self.default_conf),
+            **conf,
+        }
+        self._updata_config(detect_threshold, max_keypoints, match_threshold)
+        self._init_models()
+        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"]:
+            self.conf["matcher"]["model"]["match_threshold"] = match_threshold
+        else:
+            self.conf["feature"]["model"]["max_keypoints"] = max_keypoints
+            self.conf["feature"]["model"][
+                "keypoint_threshold"
+            ] = detect_threshold
+        self.match_conf = self.conf["matcher"]
+        self.extract_conf = self.conf["feature"]
+
+    def _init_models(self):
+        # initialize matcher
+        self.matcher = get_model(self.conf["matcher"])
+        # 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(self.match_conf["model"]["name"])
+        else:
+            pred0 = extract_features.extract(
+                self.extractor, img0, self.extract_conf["preprocessing"]
+            )
+            pred1 = extract_features.extract(
+                self.extractor, img1, self.extract_conf["preprocessing"]
+            )
+            pred = match_features.match_images(self.matcher, pred0, pred1)
+        return pred
+
+    @torch.inference_mode()
+    def 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")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    config = load_config(ROOT / "common/config.yaml")
+    test_api(config)

common/config.yaml

@@ -319,6 +319,7 @@ matcher_zoo:
       project: null
       display: true
   gluestick:
+    enable: false
     matcher: gluestick
     dense: true
     info:
@@ -329,6 +330,7 @@ matcher_zoo:
       project: https://iago-suarez.com/gluestick
       display: true
   sold2:
+    enable: false
     matcher: sold2
     dense: true
     info:

common/utils.py

@@ -78,21 +78,24 @@ def get_matcher_zoo(
     """
     matcher_zoo_restored = {}
     for k, v in matcher_zoo.items():
-        dense = v["dense"]
-        if dense:
-            matcher_zoo_restored[k] = {
-                "matcher": match_dense.confs.get(v["matcher"]),
-                "dense": dense,
-            }
-        else:
-            matcher_zoo_restored[k] = {
-                "feature": extract_features.confs.get(v["feature"]),
-                "matcher": match_features.confs.get(v["matcher"]),
-                "dense": dense,
-            }
+        matcher_zoo_restored[k] = parse_match_config(v)
     return matcher_zoo_restored
 
 
+def parse_match_config(conf):
+    if conf["dense"]:
+        return {
+            "matcher": match_dense.confs.get(conf["matcher"]),
+            "dense": True,
+        }
+    else:
+        return {
+            "feature": extract_features.confs.get(conf["feature"]),
+            "matcher": match_features.confs.get(conf["matcher"]),
+            "dense": False,
+        }
+
+
 def get_model(match_conf: Dict[str, Any]):
     """
     Load a matcher model from the provided configuration.

common/viz.py

@@ -415,12 +415,17 @@ def display_matches(
     num_inliers = 0
     KPTS0_KEY = None
     KPTS1_KEY = None
+    confid = None
     if tag == "KPTS_RAW":
         KPTS0_KEY = "mkeypoints0_orig"
         KPTS1_KEY = "mkeypoints1_orig"
+        if "mconf" in pred:
+            confid = pred["mconf"]
     elif tag == "KPTS_RANSAC":
         KPTS0_KEY = "mmkeypoints0_orig"
         KPTS1_KEY = "mmkeypoints1_orig"
+        if "mmconf" in pred:
+            confid = pred["mmconf"]
     else:
         # TODO: LINES_RAW, LINES_RANSAC
         raise ValueError(f"Unknown tag: {tag}")
@@ -434,16 +439,14 @@ def display_matches(
         mkpts0 = pred[KPTS0_KEY]
         mkpts1 = pred[KPTS1_KEY]
         num_inliers = len(mkpts0)
-        if "mmconf" in pred:
-            mmconf = pred["mmconf"]
-        else:
-            mmconf = np.ones(len(mkpts0))
+        if confid is None:
+            confid = np.ones(len(mkpts0))
         fig_mkpts = draw_matches_core(
             mkpts0,
             mkpts1,
             img0,
             img1,
-            mmconf,
+            confid,
             dpi=dpi,
             titles=titles,
             texts=texts,
@@ -472,7 +475,8 @@ def display_matches(
         # keypoints
         mkpts0 = pred.get("line_keypoints0_orig")
         mkpts1 = pred.get("line_keypoints1_orig")
-
+        fig = None
+        breakpoint()
         if mkpts0 is not None and mkpts1 is not None:
             num_inliers = len(mkpts0)
             if "mconf" in pred:

test_app_cli.py

@@ -1,153 +1,39 @@
 import cv2
 import warnings
+import numpy as np
 from pathlib import Path
 from hloc import logger
-from hloc import matchers, extractors, logger
-from hloc import match_dense, match_features, extract_features
-from hloc.utils.viz import add_text, plot_keypoints
 from common.utils import (
-    load_config,
-    get_model,
-    get_feature_model,
-    ransac_zoo,
     get_matcher_zoo,
-    filter_matches,
+    load_config,
     device,
     ROOT,
 )
-from common.viz import (
-    fig2im,
-    plot_images,
-    display_matches,
-    plot_color_line_matches,
-)
-import time
-import matplotlib.pyplot as plt
-
-warnings.simplefilter("ignore")
+from common.api import ImageMatchingAPI
 
-
-def test_modules(config: dict):
+def test_api(config: dict = None):
     img_path1 = ROOT / "datasets/sacre_coeur/mapping/02928139_3448003521.jpg"
     img_path2 = ROOT / "datasets/sacre_coeur/mapping/17295357_9106075285.jpg"
-    image0 = cv2.imread(str(img_path1))
-    image1 = cv2.imread(str(img_path2))
-    keypoint_threshold = 0.0
-    extract_max_keypoints = 2000
-    match_threshold = 0.2
-    log_path = ROOT / "experiments"
-    log_path.mkdir(exist_ok=True, parents=True)
+    image0 = cv2.imread(str(img_path1))[:, :, ::-1]
+    image1 = cv2.imread(str(img_path2))[:, :, ::-1]
 
     matcher_zoo_restored = get_matcher_zoo(config["matcher_zoo"])
     for k, v in matcher_zoo_restored.items():
         if image0 is None or image1 is None:
             logger.error("Error: No images found! Please upload two images.")
-        # init output
-        output_keypoints = None
-        output_matches_raw = None
-        output_matches_ransac = None
-        match_conf = v["matcher"]
-
-        # update match config
-        match_conf["model"]["match_threshold"] = match_threshold
-        match_conf["model"]["max_keypoints"] = extract_max_keypoints
-        matcher = get_model(match_conf)
-        t1 = time.time()
-        if v["dense"]:
-            pred = match_dense.match_images(
-                matcher,
-                image0,
-                image1,
-                match_conf["preprocessing"],
-                device=device,
-            )
-            del matcher
-            extract_conf = None
-            last_fixed = "{}".format(match_conf["model"]["name"])
+        enable = config["matcher_zoo"][k].get("enable", True)
+        if enable:
+            logger.info(f"Testing {k} ...")
+            api = ImageMatchingAPI(conf=v, device=device)
+            api(image0, image1)
+            log_path = ROOT / "experiments1"
+            log_path.mkdir(exist_ok=True, parents=True)
+            api.visualize(log_path=log_path)
         else:
-            extract_conf = v["feature"]
-
-            # update extract config
-            extract_conf["model"]["max_keypoints"] = extract_max_keypoints
-            extract_conf["model"]["keypoint_threshold"] = keypoint_threshold
-            extractor = get_feature_model(extract_conf)
-            pred0 = extract_features.extract(
-                extractor, image0, extract_conf["preprocessing"]
-            )
-            pred1 = extract_features.extract(
-                extractor, image1, extract_conf["preprocessing"]
-            )
-            pred = match_features.match_images(matcher, pred0, pred1)
-            del extractor
-            last_fixed = "{}_{}".format(
-                extract_conf["model"]["name"], match_conf["model"]["name"]
-            )
-
-        # keypoints on images
-        logger.info(f"Match features done using: {time.time()-t1:.3f}s")
-        t1 = time.time()
-        texts = [
-            f"image pairs: {img_path1.name} & {img_path2.name}",
-            "",
-        ]
-        titles = [
-            "Image 0 - Keypoints",
-            "Image 1 - Keypoints",
-        ]
-        output_keypoints = plot_images([image0, image1], titles=titles, dpi=300)
-        if "keypoints0" in pred.keys() and "keypoints1" in pred.keys():
-            plot_keypoints([pred["keypoints0"], pred["keypoints1"]])
-            text = (
-                f"# keypoints0: {len(pred['keypoints0'])} \n"
-                + f"# keypoints1: {len(pred['keypoints1'])}"
-            )
-            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
-        )
-        logger.info(f"Plot keypoints done using: {time.time()-t1:.3f}s")
-        t1 = time.time()
-
-        filter_matches(
-            pred,
-            ransac_method=config["defaults"]["ransac_method"],
-            ransac_reproj_threshold=config["defaults"][
-                "ransac_reproj_threshold"
-            ],
-            ransac_confidence=config["defaults"]["ransac_confidence"],
-            ransac_max_iter=config["defaults"]["ransac_max_iter"],
-        )
-        # 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
-        )
-        logger.info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
-
-        img_keypoints_path = log_path / f"img_keypoints_{last_fixed}.png"
-        img_matches_raw_path = log_path / f"img_matches_raw_{last_fixed}.png"
-        img_matches_ransac_path = (
-            log_path / f"img_matches_ransac_{last_fixed}.png"
-        )
-        cv2.imwrite(str(img_keypoints_path), output_keypoints)
-        cv2.imwrite(str(img_matches_raw_path), output_matches_raw)
-        cv2.imwrite(str(img_matches_ransac_path), output_matches_ransac)
-
-        plt.close("all")
-
+            logger.info(f"Skipping {k} ...")
 
 if __name__ == "__main__":
     import argparse
 
     config = load_config(ROOT / "common/config.yaml")
-    test_modules(config)
+    test_api(config)