From 318ea01f6705d1a6f6a7e6304fdc30b205a8e3c2 Mon Sep 17 00:00:00 2001
From: Realcat <realcat@126.com>
Date: Sun, 10 Nov 2024 12:08:39 +0000
Subject: [PATCH] add: xoftr

---
 .gitmodules            |  3 ++
 hloc/match_dense.py    | 19 +++++++++
 hloc/matchers/xoftr.py | 93 ++++++++++++++++++++++++++++++++++++++++++
 third_party/XoFTR      |  1 +
 ui/config.yaml         | 10 +++++
 5 files changed, 126 insertions(+)
 create mode 100644 hloc/matchers/xoftr.py
 create mode 160000 third_party/XoFTR

diff --git a/.gitmodules b/.gitmodules
index 2abaed9..43c0c1d 100644
--- a/.gitmodules
+++ b/.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
diff --git a/hloc/match_dense.py b/hloc/match_dense.py
index 29d6115..a86a62e 100644
--- a/hloc/match_dense.py
+++ b/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": {
diff --git a/hloc/matchers/xoftr.py b/hloc/matchers/xoftr.py
new file mode 100644
index 0000000..bd5f7ab
--- /dev/null
+++ b/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
diff --git a/third_party/XoFTR b/third_party/XoFTR
new file mode 160000
index 0000000..e9635d8
--- /dev/null
+++ b/third_party/XoFTR
@@ -0,0 +1 @@
+Subproject commit e9635d8baf95b5731bb1a142eff9a479a99e1e3b
diff --git a/ui/config.yaml b/ui/config.yaml
index 8f4d455..eb6b609 100644
--- a/ui/config.yaml
+++ b/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