add: xoftr

.gitmodules

@@ -70,3 +70,6 @@
 [submodule "third_party/EfficientLoFTR"]
 	path = third_party/EfficientLoFTR
 	url = https://github.com/zju3dv/EfficientLoFTR.git
+[submodule "third_party/XoFTR"]
+	path = third_party/XoFTR
+	url = https://github.com/OnderT/XoFTR.git

hloc/match_dense.py

@@ -63,6 +63,25 @@
         "max_error": 1,  # max error for assigned keypoints (in px)
         "cell_size": 1,  # size of quantization patch (max 1 kp/patch)
     },
+    "xoftr": {
+        "output": "matches-xoftr",
+        "model": {
+            "name": "xoftr",
+            "weights": "weights_xoftr_640.ckpt",
+            "max_keypoints": 2000,
+            "match_threshold": 0.3,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "width": 640,
+            "height": 480,
+            "force_resize": True,
+        },
+        "max_error": 1,  # max error for assigned keypoints (in px)
+        "cell_size": 1,  # size of quantization patch (max 1 kp/patch)
+    },
     # "loftr_quadtree": {
     #     "output": "matches-loftr-quadtree",
     #     "model": {

hloc/matchers/xoftr.py

@@ -0,0 +1,93 @@
+import sys
+import warnings
+from pathlib import Path
+
+import torch
+
+from hloc import DEVICE, MODEL_REPO_ID
+
+tp_path = Path(__file__).parent / "../../third_party"
+sys.path.append(str(tp_path))
+
+from XoFTR.src.config.default import get_cfg_defaults
+from XoFTR.src.utils.misc import lower_config
+from XoFTR.src.xoftr import XoFTR as XoFTR_
+
+from hloc import logger
+
+from ..utils.base_model import BaseModel
+
+
+class XoFTR(BaseModel):
+    default_conf = {
+        "model_name": "weights_xoftr_640.ckpt",
+        "match_threshold": 0.3,
+        "max_keypoints": -1,
+    }
+    required_inputs = ["image0", "image1"]
+
+    def _init(self, conf):
+        # Get default configurations
+        config_ = get_cfg_defaults(inference=True)
+        config_ = lower_config(config_)
+
+        # Coarse level threshold
+        config_["xoftr"]["match_coarse"]["thr"] = self.conf["match_threshold"]
+
+        # Fine level threshold
+        config_["xoftr"]["fine"]["thr"] = 0.1  # Default 0.1
+
+        # It is posseble to get denser matches
+        # If True, xoftr returns all fine-level matches for each fine-level window (at 1/2 resolution)
+        config_["xoftr"]["fine"]["denser"] = False  # Default False
+
+        # XoFTR model
+        matcher = XoFTR_(config=config_["xoftr"])
+
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+
+        # Load model
+        state_dict = torch.load(model_path, map_location="cpu")["state_dict"]
+        matcher.load_state_dict(state_dict, strict=True)
+        matcher = matcher.eval().to(DEVICE)
+        self.net = matcher
+        logger.info(f"Loaded XoFTR with weights {conf['model_name']}")
+
+    def _forward(self, data):
+        # For consistency with hloc pairs, we refine kpts in image0!
+        rename = {
+            "keypoints0": "keypoints1",
+            "keypoints1": "keypoints0",
+            "image0": "image1",
+            "image1": "image0",
+            "mask0": "mask1",
+            "mask1": "mask0",
+        }
+        data_ = {rename[k]: v for k, v in data.items()}
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            pred = self.net(data_)
+        pred = {
+            "keypoints0": data_["mkpts0_f"],
+            "keypoints1": data_["mkpts1_f"],
+        }
+        scores = data_["mconf_f"]
+
+        top_k = self.conf["max_keypoints"]
+        if top_k is not None and len(scores) > top_k:
+            keep = torch.argsort(scores, descending=True)[:top_k]
+            pred["keypoints0"], pred["keypoints1"] = (
+                pred["keypoints0"][keep],
+                pred["keypoints1"][keep],
+            )
+            scores = scores[keep]
+
+        # Switch back indices
+        pred = {(rename[k] if k in rename else k): v for k, v in pred.items()}
+        pred["scores"] = scores
+        return pred

ui/config.yaml

@@ -105,6 +105,16 @@ matcher_zoo:
       paper: https://zju3dv.github.io/efficientloftr/files/EfficientLoFTR.pdf
       project: https://zju3dv.github.io/efficientloftr
       display: true
+  xoftr:
+    matcher: xoftr
+    dense: true
+    info:
+      name: XoFTR #dispaly name
+      source: "CVPR 2024"
+      github: https://github.com/OnderT/XoFTR
+      paper: https://arxiv.org/pdf/2404.09692
+      project: null
+      display: true
   cotr:
     enable: false
     skip_ci: true