add: dkm & ense matcher max_keypoints

common/app_class.py

@@ -162,7 +162,7 @@ class ImageMatchingApp:
 
                         with gr.Accordion("Geometry Setting", open=False):
                             with gr.Row(equal_height=False):
-                                choice_estimate_geom = gr.Radio(
+                                choice_geometry_type = gr.Radio(
                                     ["Fundamental", "Homography"],
                                     label="Reconstruct Geometry",
                                     value=self.cfg["defaults"][
@@ -182,7 +182,7 @@ class ImageMatchingApp:
                         ransac_reproj_threshold,
                         ransac_confidence,
                         ransac_max_iter,
-                        choice_estimate_geom,
+                        choice_geometry_type,
                         gr.State(self.matcher_zoo),
                     ]
 
@@ -282,20 +282,20 @@ class ImageMatchingApp:
                     ransac_reproj_threshold,
                     ransac_confidence,
                     ransac_max_iter,
-                    choice_estimate_geom,
+                    choice_geometry_type,
                 ]
                 button_reset.click(
                     fn=self.ui_reset_state, inputs=None, outputs=reset_outputs
                 )
 
                 # estimate geo
-                choice_estimate_geom.change(
+                choice_geometry_type.change(
                     fn=change_estimate_geom,
                     inputs=[
                         input_image0,
                         input_image1,
                         geometry_result,
-                        choice_estimate_geom,
+                        choice_geometry_type,
                     ],
                     outputs=[output_wrapped, geometry_result],
                 )
@@ -441,12 +441,12 @@ class ImageMatchingApp:
                         v["info"]["name"],
                         v["info"]["source"],
                         v["info"]["github"],
-                        v["info"]["project"],
                         v["info"]["paper"],
+                        v["info"]["project"],
                     ]
                 )
             tab = gr.Dataframe(
-                headers=["Algo.", "Conference", "Code", "Project", "Paper"],
+                headers=["Algo.", "Conference", "Code", "Paper", "Project"],
                 datatype=["str", "str", "str", "str", "str"],
                 col_count=(5, "fixed"),
                 value=data,

common/config.yaml

@@ -1,6 +1,6 @@
 server:
   name: "0.0.0.0"
-  port: 7860
+  port: 7861
 
 defaults:
   setting_threshold: 0.1
@@ -26,6 +26,16 @@ matcher_zoo:
       paper: https://arxiv.org/abs/2305.15404
       project: https://parskatt.github.io/RoMa
       display: true
+  dkm:
+    matcher: dkm
+    dense: true
+    info: 
+      name: DKM #dispaly name
+      source: "CVPR 2023"
+      github: https://github.com/Parskatt/DKM
+      paper: https://arxiv.org/abs/2202.00667
+      project: https://parskatt.github.io/DKM
+      display: true
   loftr:
     matcher: loftr
     dense: true

common/utils.py

@@ -21,6 +21,7 @@ from .viz import (
 import time
 import matplotlib.pyplot as plt
 import warnings
+
 warnings.simplefilter("ignore")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -41,6 +42,7 @@ GRADIO_VERSION = gr.__version__.split(".")[0]
 MATCHER_ZOO = None
 models_already_loaded = {}
 
+
 def load_config(config_name: str) -> Dict[str, Any]:
     """
     Load a YAML configuration file.
@@ -417,7 +419,7 @@ def run_matching(
     ransac_reproj_threshold: int = DEFAULT_RANSAC_REPROJ_THRESHOLD,
     ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
     ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
-    choice_estimate_geom: str = DEFAULT_SETTING_GEOMETRY,
+    choice_geometry_type: str = DEFAULT_SETTING_GEOMETRY,
     matcher_zoo: Dict[str, Any] = None,
 ) -> Tuple[
     np.ndarray,
@@ -441,7 +443,7 @@ def run_matching(
         ransac_reproj_threshold (int, optional): RANSAC reprojection threshold.
         ransac_confidence (float, optional): RANSAC confidence level.
         ransac_max_iter (int, optional): RANSAC maximum number of iterations.
-        choice_estimate_geom (str, optional): setting of geometry estimation.
+        choice_geometry_type (str, optional): setting of geometry estimation.
 
     Returns:
         tuple:
@@ -476,8 +478,8 @@ def run_matching(
     cache_key = match_conf["model"]["name"]
     if cache_key in models_already_loaded:
         matcher = models_already_loaded[cache_key]
-        matcher.conf['max_keypoints'] = extract_max_keypoints
-        matcher.conf['match_threshold'] = match_threshold
+        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)
@@ -485,7 +487,7 @@ def run_matching(
     gr.Info(f"Loading model using: {time.time()-t0:.3f}s")
     logger.info(f"Loading model using: {time.time()-t0:.3f}s")
     t1 = time.time()
-    
+
     if model["dense"]:
         pred = match_dense.match_images(
             matcher, image0, image1, match_conf["preprocessing"], device=device
@@ -500,8 +502,8 @@ def run_matching(
         cache_key = extract_conf["model"]["name"]
         if cache_key in models_already_loaded:
             extractor = models_already_loaded[cache_key]
-            extractor.conf['max_keypoints'] = extract_max_keypoints
-            extractor.conf['keypoint_threshold'] = keypoint_threshold
+            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)
@@ -570,10 +572,8 @@ def run_matching(
         pred["image0_orig"],
         pred["image1_orig"],
         {"geom_info": geom_info},
-        choice_estimate_geom,
+        choice_geometry_type,
     )
-    gr.Info(f"Compute geometry done using: {time.time()-t1:.3f}s")
-    logger.info(f"Compute geometry done using: {time.time()-t1:.3f}s")
     plt.close("all")
     del pred
     logger.info(f"TOTAL time: {time.time()-t0:.3f}s")

hloc/match_dense.py

@@ -368,7 +368,7 @@ def match_images(model, image_0, image_1, conf, device="cpu"):
         }
         if "mconf" in pred.keys():
             ret["mconf"] = pred["mconf"].cpu().numpy()
-        elif "scores" in pred.keys(): #adapting loftr
+        elif "scores" in pred.keys():  # adapting loftr
             ret["mconf"] = pred["scores"].cpu().numpy()
         else:
             ret["mconf"] = np.ones_like(kpts0.cpu().numpy()[:, 0])

hloc/matchers/aspanformer.py

@@ -69,7 +69,6 @@ class ASpanFormer(BaseModel):
 
             do_system(f"cd {str(aspanformer_path)} & tar -xvf {str(tar_path)}")
 
-
         config = get_cfg_defaults()
         config.merge_from_file(conf["config_path"])
         _config = lower_config(config)
@@ -99,4 +98,14 @@ class ASpanFormer(BaseModel):
             "keypoints1": data_["mkpts1_f"],
             "mconf": data_["mconf"],
         }
+        scores = data_["mconf"]
+        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]
+            scores = scores[keep]
+            pred["keypoints0"], pred["keypoints1"], pred["mconf"] = (
+                pred["keypoints0"][keep],
+                pred["keypoints1"][keep],
+                scores,
+            )
         return pred

hloc/matchers/dkm.py

@@ -12,11 +12,13 @@ from DKM.dkm import DKMv3_outdoor
 dkm_path = Path(__file__).parent / "../../third_party/DKM"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+
 class DKMv3(BaseModel):
     default_conf = {
         "model_name": "DKMv3_outdoor.pth",
         "match_threshold": 0.2,
         "checkpoint_dir": dkm_path / "pretrained",
+        "max_keypoints": -1,
     }
     required_inputs = [
         "image0",
@@ -38,8 +40,8 @@ class DKMv3(BaseModel):
             cmd = ["wget", link, "-O", str(model_path)]
             logger.info(f"Downloading the DKMv3 model with `{cmd}`.")
             subprocess.run(cmd, check=True)
-        logger.info(f"Loading DKMv3 model...")
         self.net = DKMv3_outdoor(path_to_weights=str(model_path), device=device)
+        logger.info(f"Loading DKMv3 model done")
 
     def _forward(self, data):
         img0 = data["image0"].cpu().numpy().squeeze() * 255
@@ -52,10 +54,16 @@ class DKMv3(BaseModel):
         W_B, H_B = img1.size
 
         warp, certainty = self.net.match(img0, img1, device=device)
-        matches, certainty = self.net.sample(warp, certainty)
+        matches, certainty = self.net.sample(
+            warp, certainty, num=self.conf["max_keypoints"]
+        )
         kpts1, kpts2 = self.net.to_pixel_coordinates(
             matches, H_A, W_A, H_B, W_B
         )
-        pred = {}
-        pred["keypoints0"], pred["keypoints1"] = kpts1, kpts2
+        pred = {
+            "keypoints0": kpts1,
+            "keypoints1": kpts2,
+            "mconf": certainty,
+        }
+        breakpoint()
         return pred

hloc/matchers/loftr.py

@@ -10,15 +10,18 @@ class LoFTR(BaseModel):
     default_conf = {
         "weights": "outdoor",
         "match_threshold": 0.2,
-        "max_keypoints": None,
+        "sinkhorn_iterations": 20,
+        "max_keypoints": -1,
     }
     required_inputs = ["image0", "image1"]
 
     def _init(self, conf):
         cfg = default_cfg
         cfg["match_coarse"]["thr"] = conf["match_threshold"]
+        cfg["match_coarse"]["skh_iters"] = conf["sinkhorn_iterations"]
         self.net = LoFTR_(pretrained=conf["weights"], config=cfg)
         logger.info(f"Loaded LoFTR with weights {conf['weights']}")
+
     def _forward(self, data):
         # For consistency with hloc pairs, we refine kpts in image0!
         rename = {

hloc/matchers/topicfm.py

@@ -16,6 +16,7 @@ class TopicFM(BaseModel):
         "weights": "outdoor",
         "match_threshold": 0.2,
         "n_sampling_topics": 4,
+        "max_keypoints": -1,
     }
     required_inputs = ["image0", "image1"]
 
@@ -39,4 +40,14 @@ class TopicFM(BaseModel):
             "keypoints1": data_["mkpts1_f"],
             "mconf": data_["mconf"],
         }
+        scores = data_["mconf"]
+        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]
+            scores = scores[keep]
+            pred["keypoints0"], pred["keypoints1"], pred["mconf"] = (
+                pred["keypoints0"][keep],
+                pred["keypoints1"][keep],
+                scores,
+            )
         return pred