add: caches

common/app_class.py

@@ -34,7 +34,6 @@ class ImageMatchingApp:
         )
         self.cfg = load_config(self.config_path)
         self.matcher_zoo = get_matcher_zoo(self.cfg["matcher_zoo"])
-        # self.ransac_zoo = get_ransac_zoo(self.cfg["ransac_zoo"])
         self.app = None
         self.init_interface()
         # print all the keys

common/utils.py

@@ -39,7 +39,7 @@ DEFAULT_MATCHING_THRESHOLD = 0.2
 DEFAULT_SETTING_GEOMETRY = "Homography"
 GRADIO_VERSION = gr.__version__.split(".")[0]
 MATCHER_ZOO = None
-
+models_already_loaded = {}
 
 def load_config(config_name: str) -> Dict[str, Any]:
     """
@@ -467,15 +467,23 @@ def run_matching(
             f"Success! Please be patient and allow for about 2-3 minutes."
             f" Due to CPU inference, {key} is quiet slow."
         )
-
     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()
-    matcher = get_model(match_conf)
+    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
+        logger.info(f"Loaded cached model {cache_key}")
+    else:
+        matcher = get_model(match_conf)
+        models_already_loaded[cache_key] = matcher
     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"]:
@@ -489,7 +497,15 @@ def run_matching(
         # update extract config
         extract_conf["model"]["max_keypoints"] = extract_max_keypoints
         extract_conf["model"]["keypoint_threshold"] = keypoint_threshold
-        extractor = get_feature_model(extract_conf)
+        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
+            logger.info(f"Loaded cached model {cache_key}")
+        else:
+            extractor = get_feature_model(extract_conf)
+            models_already_loaded[cache_key] = extractor
         pred0 = extract_features.extract(
             extractor, image0, extract_conf["preprocessing"]
         )
@@ -499,6 +515,7 @@ def run_matching(
         pred = match_features.match_images(matcher, pred0, pred1)
         del extractor
     gr.Info(f"Matching images done using: {time.time()-t1:.3f}s")
+    logger.info(f"Matching images done using: {time.time()-t1:.3f}s")
     t1 = time.time()
     # plot images with keypoints
     titles = [
@@ -532,6 +549,8 @@ def run_matching(
         ransac_max_iter=ransac_max_iter,
     )
     gr.Info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
+    t1 = time.time()
 
     # plot images with ransac matches
     titles = [
@@ -541,6 +560,8 @@ def run_matching(
     output_matches_ransac, num_matches_ransac = display_matches(
         pred, titles=titles
     )
+    gr.Info(f"Display matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"Display matches done using: {time.time()-t1:.3f}s")
 
     t1 = time.time()
     # plot wrapped images
@@ -552,6 +573,7 @@ def run_matching(
         choice_estimate_geom,
     )
     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/matchers/aspanformer.py

@@ -21,6 +21,7 @@ class ASpanFormer(BaseModel):
         "weights": "outdoor",
         "match_threshold": 0.2,
         "sinkhorn_iterations": 20,
+        "max_keypoints": 2048,
         "config_path": aspanformer_path / "configs/aspan/outdoor/aspan_test.py",
         "model_name": "weights_aspanformer.tar",
     }
@@ -68,7 +69,6 @@ class ASpanFormer(BaseModel):
 
             do_system(f"cd {str(aspanformer_path)} & tar -xvf {str(tar_path)}")
 
-        logger.info(f"Loading Aspanformer model...")
 
         config = get_cfg_defaults()
         config.merge_from_file(conf["config_path"])
@@ -86,6 +86,7 @@ class ASpanFormer(BaseModel):
             "state_dict"
         ]
         self.net.load_state_dict(state_dict, strict=False)
+        logger.info(f"Loaded Aspanformer model")
 
     def _forward(self, data):
         data_ = {

hloc/matchers/loftr.py

@@ -10,7 +10,7 @@ class LoFTR(BaseModel):
     default_conf = {
         "weights": "outdoor",
         "match_threshold": 0.2,
-        "max_num_matches": None,
+        "max_keypoints": None,
     }
     required_inputs = ["image0", "image1"]
 
@@ -36,7 +36,7 @@ class LoFTR(BaseModel):
 
         scores = pred["confidence"]
 
-        top_k = self.conf["max_num_matches"]
+        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"] = (

hloc/matchers/roma.py

@@ -52,7 +52,7 @@ class Roma(BaseModel):
             logger.info(f"Downloading the dinov2 model with `{cmd}`.")
             subprocess.run(cmd, check=True)
 
-        logger.info(f"Loading Roma model...")
+        logger.info(f"Loading Roma model")
         # load the model
         weights = torch.load(model_path, map_location="cpu")
         dinov2_weights = torch.load(dinov2_weights, map_location="cpu")