Dense matching based line matcher

cfgs/localization/default.yaml

@@ -123,6 +123,9 @@ line2d:
         skip_exists: False
         superglue:
             weights: "outdoor" # ["indoor", "outdoor"] for selecting superglue models
+        dense:
+            one_to_many: False
+            weights: "outdoor" # ["indoor", "outdoor", "tiny_outdoor"] for selecting RoMa models
 var2d: # in pixels
     sold2: 5.0
     lsd: 2.0

cfgs/triangulation/default.yaml

@@ -55,6 +55,9 @@ line2d:
         skip_exists: False
         superglue:
             weights: "outdoor" # ["indoor", "outdoor"] for selecting superglue models
+        dense:
+            one_to_many: False
+            weights: "outdoor" # ["indoor", "outdoor", "tiny_outdoor"] for selecting RoMa models
 var2d: # in pixels
     sold2: 5.0
     lsd: 2.0

limap/line2d/dense/__init__.py

@@ -0,0 +1,8 @@
+from .extractor import DenseNaiveExtractor
+from .matcher import BaseDenseLineMatcherOptions, RoMaLineMatcher
+
+__all__ = [
+    "BaseDenseLineMatcherOptions",
+    "DenseNaiveExtractor",
+    "RoMaLineMatcher",
+]

limap/line2d/dense/dense_matcher/__init__.py

@@ -0,0 +1,4 @@
+from .base import BaseDenseMatcher
+from .roma import RoMa
+
+__all__ = ["BaseDenseMatcher", "RoMa"]

limap/line2d/dense/dense_matcher/base.py

@@ -0,0 +1,34 @@
+import torch
+
+
+class BaseDenseMatcher:
+    def __init__(self):
+        pass
+
+    def to_normalized_coordinates(self, coords, h, w):
+        """
+        coords: (..., 2) in the order x, y
+        """
+        coords_x = 2 / w * coords[..., 0] - 1
+        coords_y = 2 / h * coords[..., 1] - 1
+        return torch.stack([coords_x, coords_y], axis=-1)
+
+    def to_unnormalized_coordinates(self, coords, h, w):
+        """
+        Inverse operation of `to_normalized_coordinates`
+        """
+        coords_x = (coords[..., 0] + 1) * w / 2
+        coords_y = (coords[..., 1] + 1) * h / 2
+        return torch.stack([coords_x, coords_y], axis=-1)
+
+    def get_sample_thresh(self):
+        """
+        return sample threshold
+        """
+        raise NotImplementedError
+
+    def get_warping_symmetric(self, img1, img2):
+        """
+        return warp_1to2 ([-1, 1]), cert_1to2, warp_2to1([-1, 1]), cert_2to1
+        """
+        raise NotImplementedError

limap/line2d/dense/dense_matcher/roma.py

@@ -0,0 +1,48 @@
+import romatch
+from PIL import Image
+
+from .base import BaseDenseMatcher
+
+
+class RoMa(BaseDenseMatcher):
+    def __init__(self, mode="outdoor", device="cuda"):
+        super(RoMa).__init__()
+        self.output_res = 864
+        self.mode = mode
+        if mode == "outdoor":
+            self.model = romatch.roma_outdoor(
+                device=device, coarse_res=560, upsample_res=self.output_res
+            )
+        elif mode == "indoor":
+            self.model = romatch.roma_indoor(
+                device=device, coarse_res=560, upsample_res=self.output_res
+            )
+        elif mode == "tiny_outdoor":
+            self.model = romatch.tiny_roma_v1_outdoor(device=device)
+        else:
+            raise ValueError(f"Unknown mode for RoMa: {mode}")
+
+    def get_sample_thresh(self):
+        return self.model.sample_thresh
+
+    def get_warping_symmetric(self, img1, img2):
+        pil_img1 = Image.fromarray(img1)
+        pil_img2 = Image.fromarray(img2)
+        warp, certainty = self.model.match(pil_img1, pil_img2, batched=False)
+        if self.mode.startswith("tiny"):
+            warp2_to_1, certainty2_to_1 = self.model.match(
+                pil_img2, pil_img1, batched=False
+            )
+            return (
+                warp[:, :, 2:],
+                certainty,
+                warp2_to_1[:, :, 2:],
+                certainty2_to_1,
+            )
+        else:
+            return (
+                warp[:, : self.output_res, 2:],
+                certainty[:, : self.output_res],
+                warp[:, self.output_res :, :2],
+                certainty[:, self.output_res :],
+            )

limap/line2d/dense/extractor.py

@@ -0,0 +1,43 @@
+import os
+
+import numpy as np
+
+import limap.util.io as limapio
+
+from ..base_detector import BaseDetector, DefaultDetectorOptions
+
+
+class DenseNaiveExtractor(BaseDetector):
+    def __init__(self, options=DefaultDetectorOptions, device=None):
+        super().__init__(options)
+
+    def get_module_name(self):
+        return "dense_naive"
+
+    def get_descinfo_fname(self, descinfo_folder, img_id):
+        fname = os.path.join(descinfo_folder, f"descinfo_{img_id}.npz")
+        return fname
+
+    def save_descinfo(self, descinfo_folder, img_id, descinfo):
+        limapio.check_makedirs(descinfo_folder)
+        fname = self.get_descinfo_fname(descinfo_folder, img_id)
+        limapio.save_npz(fname, descinfo)
+
+    def read_descinfo(self, descinfo_folder, img_id):
+        fname = self.get_descinfo_fname(descinfo_folder, img_id)
+        descinfo = limapio.read_npz(fname)
+        return descinfo
+
+    def extract(self, camview, segs):
+        img = camview.read_image()
+        lines = segs[:, :4].reshape(-1, 2, 2)
+        scores = segs[:, -1] * np.sqrt(
+            np.linalg.norm(segs[:, :2] - segs[:, 2:4], axis=1)
+        )
+        descinfo = {
+            "image": img,
+            "image_shape": img.shape,
+            "lines": lines,
+            "scores": scores,
+        }
+        return descinfo