update: omniglue

common/utils.py

@@ -642,7 +642,7 @@ def run_matching(
     ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
     choice_geometry_type: str = DEFAULT_SETTING_GEOMETRY,
     matcher_zoo: Dict[str, Any] = None,
-    use_cached_model: bool = True,
+    use_cached_model: bool = False,
 ) -> Tuple[
     np.ndarray,
     np.ndarray,
@@ -696,19 +696,21 @@ def run_matching(
             f"Success! Please be patient and allow for about 2-3 minutes."
             f" Due to CPU inference, {key} is quiet slow."
         )
+    t0 = time.time()
     model = matcher_zoo[key]
     match_conf = model["matcher"]
     # update match config
     match_conf["model"]["match_threshold"] = match_threshold
     match_conf["model"]["max_keypoints"] = extract_max_keypoints
-    t0 = time.time()
     cache_key = "{}_{}".format(key, match_conf["model"]["name"])
-    matcher = model_cache.cache_model(cache_key, get_model, match_conf)
     if use_cached_model:
+        # because of the model cache, we need to update the config
+        matcher = model_cache.cache_model(cache_key, get_model, match_conf)
         matcher.conf["max_keypoints"] = extract_max_keypoints
         matcher.conf["match_threshold"] = match_threshold
         logger.info(f"Loaded cached model {cache_key}")
-
+    else:
+        matcher = get_model(match_conf)
     logger.info(f"Loading model using: {time.time()-t0:.3f}s")
     t1 = time.time()
 
@@ -725,13 +727,16 @@ def run_matching(
         extract_conf["model"]["keypoint_threshold"] = keypoint_threshold
         cache_key = "{}_{}".format(key, extract_conf["model"]["name"])
 
-        extractor = model_cache.cache_model(
-            cache_key, get_feature_model, extract_conf
-        )
         if use_cached_model:
+            extractor = model_cache.cache_model(
+                cache_key, get_feature_model, extract_conf
+            )
+            # because of the model cache, we need to update the config
             extractor.conf["max_keypoints"] = extract_max_keypoints
             extractor.conf["keypoint_threshold"] = keypoint_threshold
             logger.info(f"Loaded cached model {cache_key}")
+        else:
+            extractor = get_feature_model(extract_conf)
 
         pred0 = extract_features.extract(
             extractor, image0, extract_conf["preprocessing"]

hloc/match_dense.py

@@ -216,6 +216,7 @@ confs = {
         "model": {
             "name": "omniglue",
             "match_threshold": 0.2,
+            "max_keypoints": 2000,
             "features": "null",
         },
         "preprocessing": {

hloc/matchers/duster.py

@@ -105,7 +105,7 @@ class Duster(BaseModel):
         reciprocal_in_P2, nn2_in_P1, num_matches = find_reciprocal_matches(
             *pts3d_list
         )
-        print(f"found {num_matches} matches")
+        logger.info(f"Found {num_matches} matches")
         mkpts1 = pts2d_list[1][reciprocal_in_P2]
         mkpts0 = pts2d_list[0][nn2_in_P1][reciprocal_in_P2]
 
@@ -114,7 +114,6 @@ class Duster(BaseModel):
             keep = np.round(np.linspace(0, len(mkpts0) - 1, top_k)).astype(int)
             mkpts0 = mkpts0[keep]
             mkpts1 = mkpts1[keep]
-        breakpoint()
         pred = {
             "keypoints0": torch.from_numpy(mkpts0),
             "keypoints1": torch.from_numpy(mkpts1),

hloc/matchers/omniglue.py

@@ -39,7 +39,6 @@ class OmniGlue(BaseModel):
                 subprocess.run(cmd, check=True)
             else:
                 logger.error(f"Invalid dinov2 model: {dino_model_path.name}")
-
         self.net = omniglue.OmniGlue(
             og_export=str(og_model_path),
             sp_export=str(sp_model_path),
@@ -54,9 +53,8 @@ class OmniGlue(BaseModel):
         image0_rgb_np = image0_rgb_np.astype(np.uint8)  # RGB, 0-255
         image1_rgb_np = image1_rgb_np.astype(np.uint8)  # RGB, 0-255
         match_kp0, match_kp1, match_confidences = self.net.FindMatches(
-            image0_rgb_np, image1_rgb_np
+            image0_rgb_np, image1_rgb_np, self.conf["max_keypoints"]
         )
-
         # filter matches
         match_threshold = self.conf["match_threshold"]
         keep_idx = []

hloc/utils/viz.py

@@ -65,9 +65,11 @@ def plot_keypoints(kpts, colors="lime", ps=4):
     if not isinstance(colors, list):
         colors = [colors] * len(kpts)
     axes = plt.gcf().axes
-    for a, k, c in zip(axes, kpts, colors):
-        a.scatter(k[:, 0], k[:, 1], c=c, s=ps, linewidths=0)
-
+    try:
+        for a, k, c in zip(axes, kpts, colors):
+            a.scatter(k[:, 0], k[:, 1], c=c, s=ps, linewidths=0)
+    except IndexError as e:
+        pass
 
 def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.0):
     """Plot matches for a pair of existing images.

third_party/omniglue/src/omniglue/omniglue_extract.py

@@ -46,13 +46,18 @@ class OmniGlue:
                 dino_export, feature_layer=1
             )
 
-    def FindMatches(self, image0: np.ndarray, image1: np.ndarray):
+    def FindMatches(
+        self,
+        image0: np.ndarray,
+        image1: np.ndarray,
+        max_keypoints: int = 2048,
+    ):
         """TODO(omniglue): docstring."""
         height0, width0 = image0.shape[:2]
         height1, width1 = image1.shape[:2]
         # TODO: numpy to torch inputs
-        sp_features0 = self.sp_extract(image0, num_features=self.max_keypoints)
-        sp_features1 = self.sp_extract(image1, num_features=self.max_keypoints)
+        sp_features0 = self.sp_extract(image0, num_features=max_keypoints)
+        sp_features1 = self.sp_extract(image1, num_features=max_keypoints)
         dino_features0 = self.dino_extract(image0)
         dino_features1 = self.dino_extract(image1)
         dino_descriptors0 = dino_extract.get_dino_descriptors(