fix: model path

hloc/extractors/alike.py

@@ -36,13 +36,13 @@ class Alike(BaseModel):
             ),
         )
         logger.info("Loaded Alike model from {}".format(model_path))
+        configs[conf["model_name"]]["model_path"] = model_path
         self.net = Alike_(
             **configs[conf["model_name"]],
             device=device,
             top_k=conf["top_k"],
             scores_th=conf["detection_threshold"],
             n_limit=conf["max_keypoints"],
-            model_path=model_path,
         )
         logger.info("Load Alike model done.")
 

hloc/extractors/darkfeat.py

@@ -23,7 +23,7 @@ class DarkFeat(BaseModel):
     def _init(self, conf):
         model_path = self._download_model(
             repo_id=MODEL_REPO_ID,
-            filename="{}/{}.pth".format(
+            filename="{}/{}".format(
                 Path(__file__).stem, self.conf["model_name"]
             ),
         )

hloc/extractors/lanet.py

@@ -33,10 +33,6 @@ class LANet(BaseModel):
                 Path(__file__).stem, self.conf["model_name"]
             ),
         )
-        if not model_path.exists():
-            logger.warning(
-                f"No model found at {model_path}, please download it first."
-            )
         self.net = PointModel(is_test=True)
         state_dict = torch.load(model_path, map_location="cpu")
         self.net.load_state_dict(state_dict["model_state"])

hloc/extractors/sfd2.py

@@ -27,7 +27,7 @@ class SFD2(BaseModel):
         model_path = self._download_model(
             repo_id=MODEL_REPO_ID,
             filename="{}/{}".format(
-                Path(__file__).stem, self.conf["model_name"]
+                "pram", self.conf["model_name"]
             ),
         )
         self.net = load_sfd2(weight_path=model_path).eval()

hloc/match_dense.py

@@ -257,6 +257,7 @@ confs = {
         "model": {
             "name": "roma",
             "weights": "outdoor",
+            "model_name": "roma_outdoor.pth",
             "max_keypoints": 2000,
             "match_threshold": 0.2,
         },
@@ -273,7 +274,7 @@ confs = {
         "output": "matches-gim",
         "model": {
             "name": "gim",
-            "weights": "gim_dkm_100h.ckpt",
+            "model_name": "gim_dkm_100h.ckpt",
             "max_keypoints": 2000,
             "match_threshold": 0.2,
         },

hloc/matchers/gim.py

@@ -3,46 +3,116 @@ from pathlib import Path
 
 import torch
 
-from .. import MODEL_REPO_ID, logger
+from .. import MODEL_REPO_ID, logger, DEVICE
 from ..utils.base_model import BaseModel
 
 gim_path = Path(__file__).parent / "../../third_party/gim"
 sys.path.append(str(gim_path))
 
-from dkm.models.model_zoo.DKMv3 import DKMv3
+def load_model(weight_name, checkpoints_path):
+    # load model
+    model = None
+    detector = None
+    if weight_name == "gim_dkm":
+        from gim.dkm.models.model_zoo.DKMv3 import DKMv3
+        model = DKMv3(weights=None, h=672, w=896)
+    elif weight_name == "gim_loftr":
+        from gim.loftr.loftr import LoFTR
+        from gim.loftr.misc import lower_config
+        from gim.loftr.config import get_cfg_defaults
+
+        model = LoFTR(lower_config(get_cfg_defaults())["loftr"])
+    elif weight_name == "gim_lightglue":
+        from gim.lightglue.superpoint import SuperPoint
+        from gim.lightglue.models.matchers.lightglue import LightGlue
+
+        detector = SuperPoint(
+            {
+                "max_num_keypoints": 2048,
+                "force_num_keypoints": True,
+                "detection_threshold": 0.0,
+                "nms_radius": 3,
+                "trainable": False,
+            }
+        )
+        model = LightGlue(
+            {
+                "filter_threshold": 0.1,
+                "flash": False,
+                "checkpointed": True,
+            }
+        )
+
+    # load state dict
+    if weight_name == "gim_dkm":
+        state_dict = torch.load(checkpoints_path, map_location="cpu")
+        if "state_dict" in state_dict.keys():
+            state_dict = state_dict["state_dict"]
+        for k in list(state_dict.keys()):
+            if k.startswith("model."):
+                state_dict[k.replace("model.", "", 1)] = state_dict.pop(k)
+            if "encoder.net.fc" in k:
+                state_dict.pop(k)
+        model.load_state_dict(state_dict)
+
+    elif weight_name == "gim_loftr":
+        state_dict = torch.load(checkpoints_path, map_location="cpu")
+        if "state_dict" in state_dict.keys():
+            state_dict = state_dict["state_dict"]
+        model.load_state_dict(state_dict)
+
+    elif weight_name == "gim_lightglue":
+        state_dict = torch.load(checkpoints_path, map_location="cpu")
+        if "state_dict" in state_dict.keys():
+            state_dict = state_dict["state_dict"]
+        for k in list(state_dict.keys()):
+            if k.startswith("model."):
+                state_dict.pop(k)
+            if k.startswith("superpoint."):
+                state_dict[k.replace("superpoint.", "", 1)] = state_dict.pop(k)
+        detector.load_state_dict(state_dict)
+
+        state_dict = torch.load(checkpoints_path, map_location="cpu")
+        if "state_dict" in state_dict.keys():
+            state_dict = state_dict["state_dict"]
+        for k in list(state_dict.keys()):
+            if k.startswith("superpoint."):
+                state_dict.pop(k)
+            if k.startswith("model."):
+                state_dict[k.replace("model.", "", 1)] = state_dict.pop(k)
+        model.load_state_dict(state_dict)
+
+    # eval mode
+    if detector is not None:
+        detector = detector.eval().to(DEVICE)
+    model = model.eval().to(DEVICE)
+    return model
 
 
 class GIM(BaseModel):
     default_conf = {
-        "model_name": "gim_lightglue_100h.ckpt",
         "match_threshold": 0.2,
         "checkpoint_dir": gim_path / "weights",
+        "weights": "gim_dkm",
     }
     required_inputs = [
         "image0",
         "image1",
     ]
+    ckpt_name_dict = {
+        "gim_dkm": "gim_dkm_100h.ckpt",
+        "gim_loftr": "gim_loftr_50h.ckpt",
+        "gim_lightglue": "gim_lightglue_100h.ckpt",
+    }
 
     def _init(self, conf):
+        ckpt_name = self.ckpt_name_dict[conf["weights"]]
         model_path = self._download_model(
             repo_id=MODEL_REPO_ID,
-            filename="{}/{}".format(
-                Path(__file__).stem, self.conf["model_name"]
-            ),
+            filename="{}/{}".format(Path(__file__).stem, ckpt_name),
         )
-
         self.aspect_ratio = 896 / 672
-        model = DKMv3(None, 672, 896, upsample_preds=True)
-        state_dict = torch.load(str(model_path), map_location="cpu")
-        if "state_dict" in state_dict.keys():
-            state_dict = state_dict["state_dict"]
-        for k in list(state_dict.keys()):
-            if k.startswith("model."):
-                state_dict[k.replace("model.", "", 1)] = state_dict.pop(k)
-            if "encoder.net.fc" in k:
-                state_dict.pop(k)
-        model.load_state_dict(state_dict)
-
+        model = load_model(conf["weights"], model_path)
         self.net = model
         logger.info("Loaded GIM model")
 
@@ -94,6 +164,7 @@ class GIM(BaseModel):
         return mask
 
     def _forward(self, data):
+        # TODO: only support dkm+gim
         image0, image1 = self.pad_image(
             data["image0"], self.aspect_ratio
         ), self.pad_image(data["image1"], self.aspect_ratio)

hloc/matchers/imp.py

@@ -34,7 +34,7 @@ class IMP(BaseModel):
         model_path = self._download_model(
             repo_id=MODEL_REPO_ID,
             filename="{}/{}".format(
-                Path(__file__).stem, self.conf["model_name"]
+                'pram', self.conf["model_name"]
             ),
         )
 

hloc/matchers/lightglue.py

@@ -33,6 +33,7 @@ class LightGlue(BaseModel):
     ]
 
     def _init(self, conf):
+        logger.info("Loading lightglue model, {}".format(conf["model_name"]))
         model_path = self._download_model(
             repo_id=MODEL_REPO_ID,
             filename="{}/{}".format(

third_party/gim/gim/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+# @Author  : xuelun

third_party/gim/{dkm → gim/dkm}/models/dkm.py

@@ -5,9 +5,9 @@ from PIL import Image
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from dkm.utils import get_tuple_transform_ops
+from gim.dkm.utils import get_tuple_transform_ops
 from einops import rearrange
-from dkm.utils.local_correlation import local_correlation
+from gim.dkm.utils.local_correlation import local_correlation
 
 
 class ConvRefiner(nn.Module):
@@ -609,7 +609,7 @@ class RegressionMatcher(nn.Module):
         if "balanced" not in self.sample_mode:
             return good_matches, good_certainty
 
-        from dkm.utils.kde import kde
+        from gim.dkm.utils.kde import kde
         density = kde(good_matches, std=0.1)
         p = 1 / (density+1)
         p[density < 10] = 1e-7 # Basically should have at least 10 perfect neighbours, or around 100 ok ones

third_party/gim/{dkm → gim/dkm}/models/model_zoo/DKMv3.py

@@ -1,8 +1,8 @@
 import torch
 
 from torch import nn
-from dkm.models.dkm import *
-from dkm.models.encoders import *
+from gim.dkm.models.dkm import *
+from gim.dkm.models.encoders import *
 
 
 def DKMv3(weights, h, w, symmetric = True, sample_mode= "threshold_balanced", **kwargs):