add loma

Author: davnords

hloc/extract_features.py

@@ -125,6 +125,28 @@
             "resize_max": 1024,
         },
     },
+    "loma_aachen": {
+        "output": "feats-loma-n4096-r1024",
+        "model": {
+            "name": "loma",
+            "max_keypoints": 4096,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1024,
+        },
+    },
+    "loma_inloc": {
+        "output": "feats-loma-n4096-r1600",
+        "model": {
+            "name": "loma",
+            "max_keypoints": 4096,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+        },
+    },
     # Global descriptors
     "dir": {
         "output": "global-feats-dir",

hloc/extractors/loma.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn.functional as F
+from loma.descriptor.dedode import DeDoDeDescriptor
+from loma.detector.dad import DaD
+from loma.device import device
+
+from ..utils.base_model import BaseModel
+
+
+class LoMaExtractor(BaseModel):
+    default_conf = {
+        "max_keypoints": 2048,
+        "compile": False,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        # DaD weights loaded by default
+        self.detector = DaD(DaD.Cfg(compile=conf["compile"])).eval()
+
+        # Descriptor weights need to be manually loaded
+        self.descriptor = DeDoDeDescriptor(
+            DeDoDeDescriptor.Cfg(compile=conf["compile"], arch="dedode_g")
+        ).eval()
+        weights = torch.hub.load_state_dict_from_url(
+            "https://github.com/davnords/storage/releases/download/loma/loma_B.pt",
+            map_location=device,
+        )
+        weights = {k: v for k, v in weights.items() if k.startswith("_descriptor.")}
+        weights = {k[len("_descriptor.") :]: v for k, v in weights.items()}
+        self.descriptor.load_state_dict(weights, strict=True)
+
+    def preprocess_image(self, image, H=784, W=784):
+        image = F.interpolate(
+            image,
+            size=(H, W),
+            mode="bilinear",
+            align_corners=False,
+        )[0]
+        return image[None].to(device)
+
+    def detect_and_describe(self, batch: dict[str, torch.Tensor]):
+        H, W = batch["image"].shape[2:]
+
+        detections = self.detector.detect(
+            batch, num_keypoints=self.conf["max_keypoints"]
+        )
+        keypoints = detections["keypoints"]
+
+        description = self.descriptor.describe_keypoints(
+            self.preprocess_image(batch["image"]),
+            keypoints,
+        )
+        keypoints = self.detector.to_pixel_coords(keypoints, H, W)
+        keypoints = keypoints - 0.5  # be consistent with hloc
+        keypoints[..., 0] = keypoints[..., 0].clamp(0.5, W - 1.5)
+        keypoints[..., 1] = keypoints[..., 1].clamp(0.5, H - 1.5)
+        return {
+            "keypoints": [keypoints[0]],
+            "descriptors": [description["descriptions"].transpose(-1, -2)[0]],
+            "scores": [detections["keypoint_probs"][0]],
+        }
+
+    def _forward(self, data):
+        return self.detect_and_describe(data)

hloc/match_features.py

@@ -85,6 +85,10 @@
         "output": "matches-adalam",
         "model": {"name": "adalam"},
     },
+    "loma": {
+        "output": "matches-loma",
+        "model": {"name": "loma"},
+    },
 }
 
 

hloc/matchers/loma.py

@@ -0,0 +1,66 @@
+import torch
+from loma.geometry import to_normalized
+from loma.loma import LoMa, LoMaB, LoMaG, LoMaL, filter_matches
+
+from ..utils.base_model import BaseModel
+
+
+class LoMaMatcher(BaseModel):
+    default_conf = {
+        "filter_threshold": 0.1,
+        "arch": "LoMa-B",
+    }
+    required_inputs = [
+        "image0",
+        "keypoints0",
+        "descriptors0",
+        "image1",
+        "keypoints1",
+        "descriptors1",
+    ]
+
+    def _init(self, conf):
+        if conf["arch"] == "LoMa-B":
+            cfg = LoMaB()
+        elif conf["arch"] == "LoMa-L":
+            cfg = LoMaL()
+        elif conf["arch"] == "LoMa-G":
+            cfg = LoMaG()
+        else:
+            raise ValueError(f"Unknown architecture {conf['arch']}")
+        self.net = LoMa(cfg)
+
+    def _forward(self, data):
+        H_A, W_A = data["image0"].shape[2:]
+        H_B, W_B = data["image1"].shape[2:]
+
+        data["keypoints0"] = to_normalized(data["keypoints0"], H=H_A, W=W_A)
+        data["keypoints1"] = to_normalized(data["keypoints1"], H=H_B, W=W_B)
+        data["descriptors0"] = data["descriptors0"].transpose(-1, -2)
+        data["descriptors1"] = data["descriptors1"].transpose(-1, -2)
+        output = self.net(
+            data["keypoints0"],
+            data["keypoints1"],
+            data["descriptors0"],
+            data["descriptors1"],
+        )
+        scores = output["scores"]
+
+        b = data["descriptors0"].shape[0]
+        m0, m1, mscores0, mscores1 = filter_matches(
+            scores, self.conf["filter_threshold"]
+        )
+        matches, mscores = [], []
+        for k in range(b):
+            valid = m0[k] > -1
+            m_indices_0 = torch.where(valid)[0]
+            m_indices_1 = m0[k][valid]
+            matches.append(torch.stack([m_indices_0, m_indices_1], -1))
+            mscores.append(mscores0[k][valid])
+
+        return {
+            "matches0": m0,
+            "matches1": m1,
+            "matching_scores0": mscores0,
+            "matching_scores1": mscores1,
+        }

requirements.txt

@@ -11,3 +11,4 @@ pycolmap>=3.13.0
 kornia>=0.6.11
 gdown
 lightglue @ git+https://github.com/cvg/LightGlue
+lomatch @ git+https://github.com/davnords/LoMa