add: ModelCache

common/utils.py

@@ -1,7 +1,10 @@
 import os
 import cv2
+import sys
 import torch
 import random
+import psutil
+import shutil
 import numpy as np
 import gradio as gr
 from pathlib import Path
@@ -42,6 +45,66 @@ MATCHER_ZOO = None
 models_already_loaded = {}
 
 
+class ModelCache:
+    def __init__(self, max_memory_size: int = 8):
+        self.max_memory_size = max_memory_size
+        self.current_memory_size = 0
+        self.model_dict = {}
+        self.model_timestamps = []
+
+    def cache_model(self, model_key, model_loader_func, model_conf):
+        if model_key in self.model_dict:
+            self.model_timestamps.remove(model_key)
+            self.model_timestamps.append(model_key)
+            logger.info(f"Load cached {model_key}")
+            return self.model_dict[model_key]
+
+        model = self._load_model_from_disk(model_loader_func, model_conf)
+        while self._calculate_model_memory() > self.max_memory_size:
+            if len(self.model_timestamps) == 0:
+                logger.warn(
+                    "RAM: {}GB, MAX RAM: {}GB".format(
+                        self._calculate_model_memory(), self.max_memory_size
+                    )
+                )
+                break
+            oldest_model_key = self.model_timestamps.pop(0)
+            self.current_memory_size = self._calculate_model_memory()
+            logger.info(f"Del cached {oldest_model_key}")
+            del self.model_dict[oldest_model_key]
+
+        self.model_dict[model_key] = model
+        self.model_timestamps.append(model_key)
+
+        self.print_memory_usage()
+        logger.info(f"Total cached {list(self.model_dict.keys())}")
+
+        return model
+
+    def _load_model_from_disk(self, model_loader_func, model_conf):
+        return model_loader_func(model_conf)
+
+    def _calculate_model_memory(self, verbose=False):
+        host_colocation = int(os.environ.get("HOST_COLOCATION", "1"))
+        vm = psutil.virtual_memory()
+        du = shutil.disk_usage(".")
+        vm_ratio = host_colocation * vm.used / vm.total
+        if verbose:
+            logger.info(
+                f"RAM: {vm.used / 1e9:.1f}/{vm.total / host_colocation / 1e9:.1f}GB"
+            )
+            # logger.info(
+            #     f"DISK: {du.used / 1e9:.1f}/{du.total / host_colocation / 1e9:.1f}GB"
+            # )
+        return vm.used / 1e9
+
+    def print_memory_usage(self):
+        self._calculate_model_memory(verbose=True)
+
+
+model_cache = ModelCache()
+
+
 def load_config(config_name: str) -> Dict[str, Any]:
     """
     Load a YAML configuration file.
@@ -579,6 +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,
 ) -> Tuple[
     np.ndarray,
     np.ndarray,
@@ -639,15 +703,12 @@ def run_matching(
     match_conf["model"]["max_keypoints"] = extract_max_keypoints
     t0 = time.time()
     cache_key = "{}_{}".format(key, match_conf["model"]["name"])
-    if cache_key in models_already_loaded:
-        matcher = models_already_loaded[cache_key]
+    matcher = model_cache.cache_model(cache_key, get_model, match_conf)
+    if use_cached_model:
         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()
 
@@ -663,14 +724,15 @@ def run_matching(
         extract_conf["model"]["max_keypoints"] = extract_max_keypoints
         extract_conf["model"]["keypoint_threshold"] = keypoint_threshold
         cache_key = "{}_{}".format(key, extract_conf["model"]["name"])
-        if cache_key in models_already_loaded:
-            extractor = models_already_loaded[cache_key]
+
+        extractor = model_cache.cache_model(
+            cache_key, get_feature_model, extract_conf
+        )
+        if use_cached_model:
             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"]
         )

hloc/matchers/omniglue.py

@@ -10,6 +10,7 @@ from ..utils.base_model import BaseModel
 thirdparty_path = Path(__file__).parent / "../../third_party"
 sys.path.append(str(thirdparty_path))
 from omniglue.src import omniglue
+
 omniglue_path = thirdparty_path / "omniglue"
 
 

test_app_cli.py

@@ -12,11 +12,11 @@ from common.utils import (
 from common.api import ImageMatchingAPI
 
 
-def test_api(config: dict = None):
+def test_all(config: dict = None):
     img_path1 = ROOT / "datasets/sacre_coeur/mapping/02928139_3448003521.jpg"
     img_path2 = ROOT / "datasets/sacre_coeur/mapping/17295357_9106075285.jpg"
-    image0 = cv2.imread(str(img_path1))[:, :, ::-1]
-    image1 = cv2.imread(str(img_path2))[:, :, ::-1]
+    image0 = cv2.imread(str(img_path1))[:, :, ::-1]  # RGB
+    image1 = cv2.imread(str(img_path2))[:, :, ::-1]  # RGB
 
     matcher_zoo_restored = get_matcher_zoo(config["matcher_zoo"])
     for k, v in matcher_zoo_restored.items():
@@ -27,15 +27,46 @@ def test_api(config: dict = None):
             logger.info(f"Testing {k} ...")
             api = ImageMatchingAPI(conf=v, device=device)
             api(image0, image1)
-            log_path = ROOT / "experiments1"
+            log_path = ROOT / "experiments" / "all"
             log_path.mkdir(exist_ok=True, parents=True)
             api.visualize(log_path=log_path)
         else:
             logger.info(f"Skipping {k} ...")
 
 
+def test_one():
+    img_path1 = ROOT / "datasets/sacre_coeur/mapping/02928139_3448003521.jpg"
+    img_path2 = ROOT / "datasets/sacre_coeur/mapping/17295357_9106075285.jpg"
+    image0 = cv2.imread(str(img_path1))[:, :, ::-1]  # RGB
+    image1 = cv2.imread(str(img_path2))[:, :, ::-1]  # RGB
+
+    conf = {
+        "matcher": {
+            "output": "matches-omniglue",
+            "model": {
+                "name": "omniglue",
+                "match_threshold": 0.2,
+                "features": "null",
+            },
+            "preprocessing": {
+                "grayscale": False,
+                "resize_max": 1024,
+                "dfactor": 8,
+                "force_resize": False,
+            },
+        },
+        "dense": True,
+    }
+    api = ImageMatchingAPI(conf=conf, device=device)
+    api(image0, image1)
+    log_path = ROOT / "experiments" / "one"
+    log_path.mkdir(exist_ok=True, parents=True)
+    api.visualize(log_path=log_path)
+
+
 if __name__ == "__main__":
     import argparse
 
     config = load_config(ROOT / "common/config.yaml")
-    test_api(config)
+    test_one()
+    test_all(config)