update: moving lfs files

README.md

@@ -4,7 +4,7 @@ emoji: 🤗
 colorFrom: red
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.44.1
+sdk_version: 5.4.0
 app_file: app.py
 pinned: true
 license: apache-2.0

hloc/__init__.py

@@ -61,3 +61,6 @@ else:
             )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# model hub: https://huggingface.co/Realcat/imatchui_checkpoint
+MODEL_REPO_ID = "Realcat/imatchui_checkpoints"

hloc/extractors/alike.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torch
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -29,12 +29,20 @@ class Alike(BaseModel):
     required_inputs = ["image"]
 
     def _init(self, conf):
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}.pth".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        logger.info("Loaded Alike model from {}".format(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/d2net.py

@@ -1,10 +1,9 @@
-import subprocess
 import sys
 from pathlib import Path
 
 import torch
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -25,20 +24,17 @@ class D2Net(BaseModel):
     required_inputs = ["image"]
 
     def _init(self, conf):
-        model_file = conf["checkpoint_dir"] / conf["model_name"]
-        if not model_file.exists():
-            model_file.parent.mkdir(exist_ok=True)
-            cmd = [
-                "wget",
-                "--quiet",
-                "https://dusmanu.com/files/d2-net/" + conf["model_name"],
-                "-O",
-                str(model_file),
-            ]
-            subprocess.run(cmd, check=True)
 
+        logger.info("Loading D2Net model...")
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        logger.info(f"Loading model from {model_path}...")
         self.net = _D2Net(
-            model_file=model_file, use_relu=conf["use_relu"], use_cuda=False
+            model_file=model_path, use_relu=conf["use_relu"], use_cuda=False
         )
         logger.info("Load D2Net model done.")
 

hloc/extractors/darkfeat.py

@@ -2,7 +2,7 @@ import subprocess
 import sys
 from pathlib import Path
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -18,33 +18,16 @@ class DarkFeat(BaseModel):
         "detection_threshold": 0.5,
         "sub_pixel": False,
     }
-    weight_urls = {
-        "DarkFeat.pth": "https://drive.google.com/uc?id=1Thl6m8NcmQ7zSAF-1_xaFs3F4H8UU6HX&confirm=t",
-    }
-    proxy = "http://localhost:1080"
     required_inputs = ["image"]
 
     def _init(self, conf):
-        model_path = darkfeat_path / "checkpoints" / conf["model_name"]
-        link = self.weight_urls[conf["model_name"]]
-        if not model_path.exists():
-            model_path.parent.mkdir(exist_ok=True)
-            cmd_wo_proxy = ["gdown", link, "-O", str(model_path)]
-            cmd = ["gdown", link, "-O", str(model_path), "--proxy", self.proxy]
-            logger.info(
-                f"Downloading the DarkFeat model with `{cmd_wo_proxy}`."
-            )
-            try:
-                subprocess.run(cmd_wo_proxy, check=True)
-            except subprocess.CalledProcessError as e:
-                logger.info(f"Downloading the model failed `{e}`.")
-                logger.info(f"Downloading the DarkFeat model with `{cmd}`.")
-                try:
-                    subprocess.run(cmd, check=True)
-                except subprocess.CalledProcessError as e:
-                    logger.error("Failed to download the DarkFeat model.")
-                    raise e
-
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}.pth".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        logger.info("Loaded DarkFeat model: {}".format(model_path))
         self.net = DarkFeat_(model_path)
         logger.info("Load DarkFeat model done.")
 

hloc/extractors/dedode.py

@@ -1,11 +1,10 @@
-import subprocess
 import sys
 from pathlib import Path
 
 import torch
 import torchvision.transforms as transforms
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -30,39 +29,25 @@ class DeDoDe(BaseModel):
     required_inputs = [
         "image",
     ]
-    weight_urls = {
-        "dedode_detector_L.pth": "https://github.com/Parskatt/DeDoDe/releases/download/dedode_pretrained_models/dedode_detector_L.pth",
-        "dedode_descriptor_B.pth": "https://github.com/Parskatt/DeDoDe/releases/download/dedode_pretrained_models/dedode_descriptor_B.pth",
-    }
 
     # Initialize the line matcher
     def _init(self, conf):
-        model_detector_path = (
-            dedode_path / "pretrained" / conf["model_detector_name"]
+        model_detector_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, conf["model_detector_name"]
+            ),
         )
-        model_descriptor_path = (
-            dedode_path / "pretrained" / conf["model_descriptor_name"]
+        model_descriptor_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, conf["model_descriptor_name"]
+            ),
         )
-
+        logger.info("Loaded DarkFeat model: {}".format(model_detector_path))
         self.normalizer = transforms.Normalize(
             mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
         )
-        # Download the model.
-        if not model_detector_path.exists():
-            model_detector_path.parent.mkdir(exist_ok=True)
-            link = self.weight_urls[conf["model_detector_name"]]
-            cmd = ["wget", "--quiet", link, "-O", str(model_detector_path)]
-            logger.info(f"Downloading the DeDoDe detector model with `{cmd}`.")
-            subprocess.run(cmd, check=True)
-
-        if not model_descriptor_path.exists():
-            model_descriptor_path.parent.mkdir(exist_ok=True)
-            link = self.weight_urls[conf["model_descriptor_name"]]
-            cmd = ["wget", "--quiet", link, "-O", str(model_descriptor_path)]
-            logger.info(
-                f"Downloading the DeDoDe descriptor model with `{cmd}`."
-            )
-            subprocess.run(cmd, check=True)
 
         # load the model
         weights_detector = torch.load(model_detector_path, map_location="cpu")

hloc/extractors/dir.py

@@ -43,6 +43,7 @@ class DIR(BaseModel):
     }
 
     def _init(self, conf):
+        # todo: download from google drive -> huggingface models
         checkpoint = Path(
             torch.hub.get_dir(), "dirtorch", conf["model_name"] + ".pt"
         )

hloc/extractors/lanet.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torch
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -18,18 +18,25 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class LANet(BaseModel):
     default_conf = {
-        "model_name": "v0",
+        "model_name": "PointModel_v0.pth",
         "keypoint_threshold": 0.1,
         "max_keypoints": 1024,
     }
     required_inputs = ["image"]
 
     def _init(self, conf):
-        model_path = (
-            lanet_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        logger.info("Loading LANet model...")
+
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
         )
         if not model_path.exists():
-            logger.warning(f"No model found at {model_path}, start downloading")
+            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/r2d2.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torchvision.transforms as tvf
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -27,11 +27,16 @@ class R2D2(BaseModel):
     required_inputs = ["image"]
 
     def _init(self, conf):
-        model_fn = r2d2_path / "models" / conf["model_name"]
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
         self.norm_rgb = tvf.Normalize(
             mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
         )
-        self.net = load_network(model_fn)
+        self.net = load_network(model_path)
         self.detector = NonMaxSuppression(
             rel_thr=conf["reliability_threshold"],
             rep_thr=conf["repetability_threshold"],

hloc/extractors/rekd.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torch
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -22,8 +22,12 @@ class REKD(BaseModel):
     required_inputs = ["image"]
 
     def _init(self, conf):
-        model_path = (
-            rekd_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        # TODO: download model
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
         )
         if not model_path.exists():
             print(f"No model found at {model_path}")

hloc/extractors/rord.py

@@ -4,7 +4,7 @@ from pathlib import Path
 
 import torch
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 
 from ..utils.base_model import BaseModel
 
@@ -23,28 +23,14 @@ class RoRD(BaseModel):
         "max_keypoints": 1024,
     }
     required_inputs = ["image"]
-    weight_urls = {
-        "rord.pth": "https://drive.google.com/uc?id=12414ZGKwgPAjNTGtNrlB4VV9l7W76B2o&confirm=t",
-    }
-    proxy = "http://localhost:1080"
 
     def _init(self, conf):
-        model_path = conf["checkpoint_dir"] / conf["model_name"]
-        link = self.weight_urls[conf["model_name"]]
-        if not model_path.exists():
-            model_path.parent.mkdir(exist_ok=True)
-            cmd_wo_proxy = ["gdown", link, "-O", str(model_path)]
-            cmd = ["gdown", link, "-O", str(model_path), "--proxy", self.proxy]
-            logger.info(f"Downloading the RoRD model with `{cmd_wo_proxy}`.")
-            try:
-                subprocess.run(cmd_wo_proxy, check=True)
-            except subprocess.CalledProcessError as e:
-                logger.info(f"Downloading failed {e}.")
-                logger.info(f"Downloading the RoRD model with {cmd}.")
-                try:
-                    subprocess.run(cmd, check=True)
-                except subprocess.CalledProcessError as e:
-                    logger.error(f"Failed to download the RoRD model: {e}")
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
         self.net = _RoRD(
             model_file=model_path, use_relu=conf["use_relu"], use_cuda=False
         )

hloc/extractors/sfd2.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torchvision.transforms as tvf
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 tp_path = Path(__file__).parent / "../../third_party"
@@ -24,7 +24,12 @@ class SFD2(BaseModel):
         self.norm_rgb = tvf.Normalize(
             mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
         )
-        model_path = tp_path / "pram" / "weights" / self.conf["model_name"]
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
         self.net = load_sfd2(weight_path=model_path).eval()
 
         logger.info("Load SFD2 model done.")

hloc/matchers/aspanformer.py

@@ -4,7 +4,7 @@ from pathlib import Path
 
 import torch
 
-from hloc import logger
+from hloc import MODEL_REPO_ID, logger
 from hloc.utils.base_model import BaseModel
 
 sys.path.append(str(Path(__file__).parent / "../../third_party"))
@@ -17,59 +17,15 @@ aspanformer_path = Path(__file__).parent / "../../third_party/ASpanFormer"
 
 class ASpanFormer(BaseModel):
     default_conf = {
-        "weights": "outdoor",
+        "model_name": "outdoor.ckpt",
         "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",
     }
     required_inputs = ["image0", "image1"]
-    proxy = "http://localhost:1080"
-    aspanformer_models = {
-        "weights_aspanformer.tar": "https://drive.google.com/uc?id=1eavM9dTkw9nbc-JqlVVfGPU5UvTTfc6k&confirm=t"
-    }
 
     def _init(self, conf):
-        model_path = (
-            aspanformer_path / "weights" / Path(conf["weights"] + ".ckpt")
-        )
-        # Download the model.
-        if not model_path.exists():
-            # model_path.parent.mkdir(exist_ok=True)
-            tar_path = aspanformer_path / conf["model_name"]
-            if not tar_path.exists():
-                link = self.aspanformer_models[conf["model_name"]]
-                cmd = [
-                    "gdown",
-                    link,
-                    "-O",
-                    str(tar_path),
-                    "--proxy",
-                    self.proxy,
-                ]
-                cmd_wo_proxy = ["gdown", link, "-O", str(tar_path)]
-                logger.info(
-                    f"Downloading the Aspanformer model with `{cmd_wo_proxy}`."
-                )
-                try:
-                    subprocess.run(cmd_wo_proxy, check=True)
-                except subprocess.CalledProcessError as e:
-                    logger.info(f"Downloading failed {e}.")
-                    logger.info(
-                        f"Downloading the Aspanformer model with `{cmd}`."
-                    )
-                    try:
-                        subprocess.run(cmd, check=True)
-                    except subprocess.CalledProcessError as e:
-                        logger.error(
-                            f"Failed to download the Aspanformer model: {e}"
-                        )
-
-            cmd = ["tar", "-xvf", str(tar_path), "-C", str(aspanformer_path)]
-            logger.info(f"Unzip model file `{cmd}`.")
-            subprocess.run(cmd, check=True)
-
         config = get_cfg_defaults()
         config.merge_from_file(conf["config_path"])
         _config = lower_config(config)
@@ -81,8 +37,13 @@ class ASpanFormer(BaseModel):
         ]
 
         self.net = _ASpanFormer(config=_config["aspan"])
-        weight_path = model_path
-        state_dict = torch.load(str(weight_path), map_location="cpu")[
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        state_dict = torch.load(str(model_path), map_location="cpu")[
             "state_dict"
         ]
         self.net.load_state_dict(state_dict, strict=False)

hloc/matchers/cotr.py

@@ -6,6 +6,8 @@ import numpy as np
 import torch
 from torchvision.transforms import ToPILImage
 
+from hloc import DEVICE, MODEL_REPO_ID
+
 from ..utils.base_model import BaseModel
 
 sys.path.append(str(Path(__file__).parent / "../../third_party/COTR"))
@@ -18,16 +20,13 @@ from COTR.utils import utils as utils_cotr
 utils_cotr.fix_randomness(0)
 torch.set_grad_enabled(False)
 
-cotr_path = Path(__file__).parent / "../../third_party/COTR"
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 class COTR(BaseModel):
     default_conf = {
         "weights": "out/default",
         "match_threshold": 0.2,
         "max_keypoints": -1,
+        "model_name": "checkpoint.pth.tar",
     }
     required_inputs = ["image0", "image1"]
 
@@ -36,8 +35,11 @@ class COTR(BaseModel):
         set_COTR_arguments(parser)  # noqa: F405
         opt = parser.parse_args()
         opt.command = " ".join(sys.argv)
-        opt.load_weights_path = str(
-            cotr_path / conf["weights"] / "checkpoint.pth.tar"
+        opt.load_weights_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
         )
 
         layer_2_channels = {
@@ -49,7 +51,7 @@ class COTR(BaseModel):
         opt.dim_feedforward = layer_2_channels[opt.layer]
 
         model = build_model(opt)
-        model = model.to(device)
+        model = model.to(DEVICE)
         weights = torch.load(opt.load_weights_path, map_location="cpu")[
             "model_state_dict"
         ]

hloc/matchers/dkm.py

@@ -1,48 +1,36 @@
-import subprocess
 import sys
 from pathlib import Path
 
 import torch
 from PIL import Image
 
-from .. import logger
-from ..utils.base_model import BaseModel
+from hloc import DEVICE, MODEL_REPO_ID, logger
+from hloc.utils.base_model import BaseModel
 
 sys.path.append(str(Path(__file__).parent / "../../third_party"))
 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",
         "image1",
     ]
-    # Models exported using
-    dkm_models = {
-        "DKMv3_outdoor.pth": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_outdoor.pth",
-        "DKMv3_indoor.pth": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_indoor.pth",
-    }
 
     def _init(self, conf):
-        model_path = dkm_path / "pretrained" / conf["model_name"]
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
 
-        # Download the model.
-        if not model_path.exists():
-            model_path.parent.mkdir(exist_ok=True)
-            link = self.dkm_models[conf["model_name"]]
-            cmd = ["wget", "--quiet", link, "-O", str(model_path)]
-            logger.info(f"Downloading the DKMv3 model with `{cmd}`.")
-            subprocess.run(cmd, check=True)
-        self.net = DKMv3_outdoor(path_to_weights=str(model_path), device=device)
+        self.net = DKMv3_outdoor(path_to_weights=str(model_path), device=DEVICE)
         logger.info("Loading DKMv3 model done")
 
     def _forward(self, data):
@@ -55,7 +43,7 @@ class DKMv3(BaseModel):
         W_A, H_A = img0.size
         W_B, H_B = img1.size
 
-        warp, certainty = self.net.match(img0, img1, device=device)
+        warp, certainty = self.net.match(img0, img1, device=DEVICE)
         matches, certainty = self.net.sample(
             warp, certainty, num=self.conf["max_keypoints"]
         )

hloc/matchers/duster.py

@@ -1,13 +1,11 @@
-import os
 import sys
-import urllib.request
 from pathlib import Path
 
 import numpy as np
 import torch
 import torchvision.transforms as tfm
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 duster_path = Path(__file__).parent / "../../third_party/dust3r"
@@ -25,30 +23,24 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 class Duster(BaseModel):
     default_conf = {
         "name": "Duster3r",
-        "model_path": duster_path / "model_weights/duster_vit_large.pth",
+        "model_name": "duster_vit_large.pth",
         "max_keypoints": 3000,
         "vit_patch_size": 16,
     }
 
     def _init(self, conf):
         self.normalize = tfm.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-        self.model_path = self.conf["model_path"]
-        self.download_weights()
-        # self.net = load_model(self.model_path, device)
-        self.net = AsymmetricCroCo3DStereo.from_pretrained(
-            self.model_path
-            # "naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt"
-        ).to(device)
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        self.net = AsymmetricCroCo3DStereo.from_pretrained(model_path).to(
+            device
+        )
         logger.info("Loaded Dust3r model")
 
-    def download_weights(self):
-        url = "https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth"
-
-        self.model_path.parent.mkdir(parents=True, exist_ok=True)
-        if not os.path.isfile(self.model_path):
-            logger.info("Downloading Duster(ViT large)... (takes a while)")
-            urllib.request.urlretrieve(url, self.model_path)
-
     def preprocess(self, img):
         # the super-class already makes sure that img0,img1 have
         # same resolution and that h == w

hloc/matchers/eloftr.py

@@ -5,6 +5,8 @@ from pathlib import Path
 
 import torch
 
+from hloc import MODEL_REPO_ID
+
 tp_path = Path(__file__).parent / "../../third_party"
 sys.path.append(str(tp_path))
 
@@ -22,7 +24,7 @@ from ..utils.base_model import BaseModel
 
 class ELoFTR(BaseModel):
     default_conf = {
-        "weights": "weights/eloftr_outdoor.ckpt",
+        "model_name": "eloftr_outdoor.ckpt",
         "match_threshold": 0.2,
         # "sinkhorn_iterations": 20,
         "max_keypoints": -1,
@@ -44,7 +46,14 @@ class ELoFTR(BaseModel):
             _default_cfg["mp"] = True
         elif self.conf["precision"] == "fp16":
             _default_cfg["half"] = True
-        model_path = tp_path / "EfficientLoFTR" / self.conf["weights"]
+
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+
         cfg = _default_cfg
         cfg["match_coarse"]["thr"] = conf["match_threshold"]
         # cfg["match_coarse"]["skh_iters"] = conf["sinkhorn_iterations"]
@@ -55,7 +64,7 @@ class ELoFTR(BaseModel):
 
         if self.conf["precision"] == "fp16":
             self.net = self.net.half()
-        logger.info(f"Loaded Efficient LoFTR with weights {conf['weights']}")
+        logger.info(f"Loaded Efficient LoFTR with weights {conf['model_name']}")
 
     def _forward(self, data):
         # For consistency with hloc pairs, we refine kpts in image0!

hloc/matchers/gim.py

@@ -1,11 +1,9 @@
-import subprocess
 import sys
 from pathlib import Path
 
-import gdown
 import torch
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 gim_path = Path(__file__).parent / "../../third_party/gim"
@@ -13,12 +11,10 @@ sys.path.append(str(gim_path))
 
 from dkm.models.model_zoo.DKMv3 import DKMv3
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 class GIM(BaseModel):
     default_conf = {
-        "model_name": "gim_dkm_100h.ckpt",
+        "model_name": "gim_lightglue_100h.ckpt",
         "match_threshold": 0.2,
         "checkpoint_dir": gim_path / "weights",
     }
@@ -26,27 +22,14 @@ class GIM(BaseModel):
         "image0",
         "image1",
     ]
-    model_dict = {
-        "gim_lightglue_100h.ckpt": "https://github.com/xuelunshen/gim/blob/main/weights/gim_lightglue_100h.ckpt",
-        "gim_dkm_100h.ckpt": "https://drive.google.com/file/d/1gk97V4IROnR1Nprq10W9NCFUv2mxXR_-/view",
-    }
 
     def _init(self, conf):
-        conf["model_name"] = str(conf["weights"])
-        if conf["model_name"] not in self.model_dict:
-            raise ValueError(f"Unknown GIM model {conf['model_name']}.")
-        model_path = conf["checkpoint_dir"] / conf["model_name"]
-
-        # Download the model.
-        if not model_path.exists():
-            model_path.parent.mkdir(exist_ok=True)
-            model_link = self.model_dict[conf["model_name"]]
-            if "drive.google.com" in model_link:
-                gdown.download(model_link, output=str(model_path), fuzzy=True)
-            else:
-                cmd = ["wget", "--quiet", model_link, "-O", str(model_path)]
-                subprocess.run(cmd, check=True)
-            logger.info("Downloaded GIM model succeeed!")
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
 
         self.aspect_ratio = 896 / 672
         model = DKMv3(None, 672, 896, upsample_preds=True)

hloc/matchers/gluestick.py

@@ -1,10 +1,9 @@
-import subprocess
 import sys
 from pathlib import Path
 
 import torch
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 gluestick_path = Path(__file__).parent / "../../third_party/GlueStick"
@@ -13,8 +12,6 @@ sys.path.append(str(gluestick_path))
 from gluestick import batch_to_np
 from gluestick.models.two_view_pipeline import TwoViewPipeline
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 class GlueStick(BaseModel):
     default_conf = {
@@ -30,23 +27,15 @@ class GlueStick(BaseModel):
         "image1",
     ]
 
-    gluestick_models = {
-        "checkpoint_GlueStick_MD.tar": "https://github.com/cvg/GlueStick/releases/download/v0.1_arxiv/checkpoint_GlueStick_MD.tar",
-    }
-
     # Initialize the line matcher
     def _init(self, conf):
-        model_path = (
-            gluestick_path / "resources" / "weights" / conf["model_name"]
-        )
-
         # Download the model.
-        if not model_path.exists():
-            model_path.parent.mkdir(exist_ok=True)
-            link = self.gluestick_models[conf["model_name"]]
-            cmd = ["wget", "--quiet", link, "-O", str(model_path)]
-            logger.info(f"Downloading the Gluestick model with `{cmd}`.")
-            subprocess.run(cmd, check=True)
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
         logger.info("Loading GlueStick model...")
 
         gluestick_conf = {

hloc/matchers/imp.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torch
 
-from .. import DEVICE, logger
+from .. import DEVICE, MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 tp_path = Path(__file__).parent / "../../third_party"
@@ -31,11 +31,17 @@ class IMP(BaseModel):
 
     def _init(self, conf):
         self.conf = {**self.default_conf, **conf}
-        weight_path = tp_path / "pram" / "weights" / self.conf["model_name"]
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+
         # self.net = nets.gml(self.conf).eval().to(DEVICE)
         self.net = GML(self.conf).eval().to(DEVICE)
         self.net.load_state_dict(
-            torch.load(weight_path, map_location="cpu")["model"], strict=True
+            torch.load(model_path, map_location="cpu")["model"], strict=True
         )
         logger.info("Load IMP model done.")
 

hloc/matchers/lightglue.py

@@ -1,7 +1,7 @@
 import sys
 from pathlib import Path
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 lightglue_path = Path(__file__).parent / "../../third_party/LightGlue"
@@ -33,8 +33,13 @@ class LightGlue(BaseModel):
     ]
 
     def _init(self, conf):
-        weight_path = lightglue_path / "weights" / conf["model_name"]
-        conf["weights"] = str(weight_path)
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        conf["weights"] = str(model_path)
         conf["filter_threshold"] = conf["match_threshold"]
         self.net = LG(**conf)
         logger.info("Load lightglue model done.")

hloc/matchers/mast3r.py

@@ -1,13 +1,11 @@
-import os
 import sys
-import urllib.request
 from pathlib import Path
 
 import numpy as np
 import torch
 import torchvision.transforms as tfm
 
-from .. import logger
+from .. import DEVICE, MODEL_REPO_ID, logger
 
 mast3r_path = Path(__file__).parent / "../../third_party/mast3r"
 sys.path.append(str(mast3r_path))
@@ -22,38 +20,30 @@ from mast3r.model import AsymmetricMASt3R
 
 from hloc.matchers.duster import Duster
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 class Mast3r(Duster):
     default_conf = {
         "name": "Mast3r",
-        "model_path": mast3r_path
-        / "model_weights/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth",
+        "model_name": "MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth",
         "max_keypoints": 2000,
         "vit_patch_size": 16,
     }
 
     def _init(self, conf):
         self.normalize = tfm.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-        self.model_path = self.conf["model_path"]
-        self.download_weights()
-        self.net = AsymmetricMASt3R.from_pretrained(self.model_path).to(device)
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        self.net = AsymmetricMASt3R.from_pretrained(model_path).to(DEVICE)
         logger.info("Loaded Mast3r model")
 
-    def download_weights(self):
-        url = "https://download.europe.naverlabs.com/ComputerVision/MASt3R/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth"
-
-        self.model_path.parent.mkdir(parents=True, exist_ok=True)
-        if not os.path.isfile(self.model_path):
-            logger.info("Downloading Mast3r(ViT large)... (takes a while)")
-            urllib.request.urlretrieve(url, self.model_path)
-            logger.info("Downloading Mast3r(ViT large)... done!")
-
     def _forward(self, data):
         img0, img1 = data["image0"], data["image1"]
-        mean = torch.tensor([0.5, 0.5, 0.5]).to(device)
-        std = torch.tensor([0.5, 0.5, 0.5]).to(device)
+        mean = torch.tensor([0.5, 0.5, 0.5]).to(DEVICE)
+        std = torch.tensor([0.5, 0.5, 0.5]).to(DEVICE)
 
         img0 = (img0 - mean.view(1, 3, 1, 1)) / std.view(1, 3, 1, 1)
         img1 = (img1 - mean.view(1, 3, 1, 1)) / std.view(1, 3, 1, 1)
@@ -65,7 +55,7 @@ class Mast3r(Duster):
         pairs = make_pairs(
             images, scene_graph="complete", prefilter=None, symmetrize=True
         )
-        output = inference(pairs, self.net, device, batch_size=1)
+        output = inference(pairs, self.net, DEVICE, batch_size=1)
 
         # at this stage, you have the raw dust3r predictions
         _, pred1 = output["view1"], output["pred1"]
@@ -81,7 +71,7 @@ class Mast3r(Duster):
             desc1,
             desc2,
             subsample_or_initxy1=2,
-            device=device,
+            device=DEVICE,
             dist="dot",
             block_size=2**13,
         )

hloc/matchers/mickey.py

@@ -1,10 +1,9 @@
-import subprocess
 import sys
 from pathlib import Path
 
 import torch
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 mickey_path = Path(__file__).parent / "../../third_party"
@@ -13,8 +12,6 @@ sys.path.append(str(mickey_path))
 from mickey.config.default import cfg
 from mickey.lib.models.builder import build_model
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 class Mickey(BaseModel):
     default_conf = {
@@ -26,33 +23,23 @@ class Mickey(BaseModel):
         "image0",
         "image1",
     ]
-    weight_urls = "https://storage.googleapis.com/niantic-lon-static/research/mickey/assets/mickey_weights.zip"
 
     # Initialize the line matcher
     def _init(self, conf):
-        model_path = mickey_path / "mickey/mickey_weights" / conf["model_name"]
-        zip_path = mickey_path / "mickey/mickey_weights.zip"
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        # TODO: config path of mickey
         config_path = model_path.parent / self.conf["config_path"]
-        # Download the model.
-        if not model_path.exists():
-            model_path.parent.mkdir(exist_ok=True, parents=True)
-            link = self.weight_urls
-            if not zip_path.exists():
-                cmd = ["wget", "--quiet", link, "-O", str(zip_path)]
-                logger.info(f"Downloading the Mickey model with {cmd}.")
-                subprocess.run(cmd, check=True)
-            cmd = ["unzip", "-d", str(model_path.parent.parent), str(zip_path)]
-            logger.info(f"Running {cmd}.")
-            subprocess.run(cmd, check=True)
-
         logger.info("Loading mickey model...")
         cfg.merge_from_file(config_path)
         self.net = build_model(cfg, checkpoint=model_path)
         logger.info("Load Mickey model done.")
 
     def _forward(self, data):
-        # data['K_color0'] = torch.from_numpy(K['im0.jpg']).unsqueeze(0).to(device)
-        # data['K_color1'] = torch.from_numpy(K['im1.jpg']).unsqueeze(0).to(device)
         pred = self.net(data)
         pred = {
             **pred,

hloc/matchers/omniglue.py

@@ -5,7 +5,7 @@ from pathlib import Path
 import numpy as np
 import torch
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 thirdparty_path = Path(__file__).parent / "../../third_party"
@@ -27,19 +27,21 @@ class OmniGlue(BaseModel):
 
     def _init(self, conf):
         logger.info("Loading OmniGlue model")
-        og_model_path = omniglue_path / "models" / "omniglue.onnx"
-        sp_model_path = omniglue_path / "models" / "sp_v6.onnx"
-        dino_model_path = (
-            omniglue_path / "models" / "dinov2_vitb14_pretrain.pth"  # ~330MB
+        og_model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, "omniglue.onnx"),
         )
-        if not dino_model_path.exists():
-            link = self.dino_v2_link_dict.get(dino_model_path.name, None)
-            if link is not None:
-                cmd = ["wget", "--quiet", link, "-O", str(dino_model_path)]
-                logger.info(f"Downloading the dinov2 model with `{cmd}`.")
-                subprocess.run(cmd, check=True)
-            else:
-                logger.error(f"Invalid dinov2 model: {dino_model_path.name}")
+        sp_model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, "sp_v6.onnx"),
+        )
+        dino_model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, "dinov2_vitb14_pretrain.pth"
+            ),
+        )
+
         self.net = omniglue.OmniGlue(
             og_export=str(og_model_path),
             sp_export=str(sp_model_path),

hloc/matchers/roma.py

@@ -1,11 +1,10 @@
-import subprocess
 import sys
 from pathlib import Path
 
 import torch
 from PIL import Image
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 roma_path = Path(__file__).parent / "../../third_party/RoMa"
@@ -26,33 +25,22 @@ class Roma(BaseModel):
         "image0",
         "image1",
     ]
-    weight_urls = {
-        "roma": {
-            "roma_outdoor.pth": "https://github.com/Parskatt/storage/releases/download/roma/roma_outdoor.pth",
-            "roma_indoor.pth": "https://github.com/Parskatt/storage/releases/download/roma/roma_indoor.pth",
-        },
-        "dinov2_vitl14_pretrain.pth": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth",
-    }
 
     # Initialize the line matcher
     def _init(self, conf):
-        model_path = roma_path / "pretrained" / conf["model_name"]
-        dinov2_weights = roma_path / "pretrained" / conf["model_utils_name"]
-
-        # Download the model.
-        if not model_path.exists():
-            model_path.parent.mkdir(exist_ok=True)
-            link = self.weight_urls["roma"][conf["model_name"]]
-            cmd = ["wget", "--quiet", link, "-O", str(model_path)]
-            logger.info(f"Downloading the Roma model with `{cmd}`.")
-            subprocess.run(cmd, check=True)
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
 
-        if not dinov2_weights.exists():
-            dinov2_weights.parent.mkdir(exist_ok=True)
-            link = self.weight_urls[conf["model_utils_name"]]
-            cmd = ["wget", "--quiet", link, "-O", str(dinov2_weights)]
-            logger.info(f"Downloading the dinov2 model with `{cmd}`.")
-            subprocess.run(cmd, check=True)
+        dinov2_weights = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_utils_name"]
+            ),
+        )
 
         logger.info("Loading Roma model")
         # load the model

hloc/matchers/sgmnet.py

@@ -1,11 +1,10 @@
-import subprocess
 import sys
 from collections import OrderedDict, namedtuple
 from pathlib import Path
 
 import torch
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 sgmnet_path = Path(__file__).parent / "../../third_party/SGMNet"
@@ -19,7 +18,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 class SGMNet(BaseModel):
     default_conf = {
         "name": "SGM",
-        "model_name": "model_best.pth",
+        "model_name": "weights/sgm/root/model_best.pth",
         "seed_top_k": [256, 256],
         "seed_radius_coe": 0.01,
         "net_channels": 128,
@@ -37,50 +36,20 @@ class SGMNet(BaseModel):
         "image0",
         "image1",
     ]
-    weight_urls = {
-        "model_best.pth": "https://drive.google.com/uc?id=1Ca0WmKSSt2G6P7m8YAOlSAHEFar_TAWb&confirm=t",
-    }
-    proxy = "http://localhost:1080"
 
     # Initialize the line matcher
     def _init(self, conf):
-        sgmnet_weights = sgmnet_path / "weights/sgm/root" / conf["model_name"]
-
-        link = self.weight_urls[conf["model_name"]]
-        tar_path = sgmnet_path / "weights.tar.gz"
-        # Download the model.
-        if not sgmnet_weights.exists():
-            if not tar_path.exists():
-                cmd = [
-                    "gdown",
-                    link,
-                    "-O",
-                    str(tar_path),
-                    "--proxy",
-                    self.proxy,
-                ]
-                cmd_wo_proxy = ["gdown", link, "-O", str(tar_path)]
-                logger.info(
-                    f"Downloading the SGMNet model with `{cmd_wo_proxy}`."
-                )
-                try:
-                    subprocess.run(cmd_wo_proxy, check=True)
-                except subprocess.CalledProcessError as e:
-                    logger.info(f"Downloading failed {e}.")
-                    logger.info(f"Downloading the SGMNet model with `{cmd}`.")
-                    try:
-                        subprocess.run(cmd, check=True)
-                    except subprocess.CalledProcessError as e:
-                        logger.error("Failed to download the SGMNet model.")
-                        raise e
-            cmd = ["tar", "-xvf", str(tar_path), "-C", str(sgmnet_path)]
-            logger.info(f"Unzip model file `{cmd}`.")
-            subprocess.run(cmd, check=True)
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
 
         # config
         config = namedtuple("config", conf.keys())(*conf.values())
         self.net = SGM_Model(config)
-        checkpoint = torch.load(sgmnet_weights, map_location="cpu")
+        checkpoint = torch.load(model_path, map_location="cpu")
         # for ddp model
         if (
             list(checkpoint["state_dict"].items())[0][0].split(".")[0]

hloc/matchers/sold2.py

@@ -1,10 +1,9 @@
-import subprocess
 import sys
 from pathlib import Path
 
 import torch
 
-from .. import logger
+from .. import MODEL_REPO_ID, logger
 from ..utils.base_model import BaseModel
 
 sold2_path = Path(__file__).parent / "../../third_party/SOLD2"
@@ -17,7 +16,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class SOLD2(BaseModel):
     default_conf = {
-        "weights": "sold2_wireframe.tar",
+        "model_name": "sold2_wireframe.tar",
         "match_threshold": 0.2,
         "checkpoint_dir": sold2_path / "pretrained",
         "detect_thresh": 0.25,
@@ -31,21 +30,15 @@ class SOLD2(BaseModel):
         "image1",
     ]
 
-    weight_urls = {
-        "sold2_wireframe.tar": "https://www.polybox.ethz.ch/index.php/s/blOrW89gqSLoHOk/download",
-    }
-
     # Initialize the line matcher
     def _init(self, conf):
-        checkpoint_path = conf["checkpoint_dir"] / conf["weights"]
-
-        # Download the model.
-        if not checkpoint_path.exists():
-            checkpoint_path.parent.mkdir(exist_ok=True)
-            link = self.weight_urls[conf["weights"]]
-            cmd = ["wget", "--quiet", link, "-O", str(checkpoint_path)]
-            logger.info(f"Downloading the SOLD2 model with `{cmd}`.")
-            subprocess.run(cmd, check=True)
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        logger.info("Loading SOLD2 model: {}".format(model_path))
 
         mode = "dynamic"  # 'dynamic' or 'static'
         match_config = {
@@ -127,7 +120,7 @@ class SOLD2(BaseModel):
         }
         self.net = LineMatcher(
             match_config["model_cfg"],
-            checkpoint_path,
+            model_path,
             device,
             match_config["line_detector_cfg"],
             match_config["line_matcher_cfg"],

hloc/matchers/topicfm.py

@@ -3,6 +3,8 @@ from pathlib import Path
 
 import torch
 
+from hloc import MODEL_REPO_ID
+
 from ..utils.base_model import BaseModel
 
 sys.path.append(str(Path(__file__).parent / "../../third_party"))
@@ -15,6 +17,7 @@ topicfm_path = Path(__file__).parent / "../../third_party/TopicFM"
 class TopicFM(BaseModel):
     default_conf = {
         "weights": "outdoor",
+        "model_name": "model_best.ckpt",
         "match_threshold": 0.2,
         "n_sampling_topics": 4,
         "max_keypoints": -1,
@@ -25,9 +28,14 @@ class TopicFM(BaseModel):
         _conf = dict(get_model_cfg())
         _conf["match_coarse"]["thr"] = conf["match_threshold"]
         _conf["coarse"]["n_samples"] = conf["n_sampling_topics"]
-        weight_path = topicfm_path / "pretrained/model_best.ckpt"
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
         self.net = _TopicFM(config=_conf)
-        ckpt_dict = torch.load(weight_path, map_location="cpu")
+        ckpt_dict = torch.load(model_path, map_location="cpu")
         self.net.load_state_dict(ckpt_dict["state_dict"])
 
     def _forward(self, data):

hloc/utils/base_model.py

@@ -3,6 +3,7 @@ from abc import ABCMeta, abstractmethod
 from torch import nn
 from copy import copy
 import inspect
+from huggingface_hub import hf_hub_download
 
 
 class BaseModel(nn.Module, metaclass=ABCMeta):
@@ -32,7 +33,14 @@ class BaseModel(nn.Module, metaclass=ABCMeta):
     def _forward(self, data):
         """To be implemented by the child class."""
         raise NotImplementedError
-
+    
+    def _download_model(self, repo_id=None, filename=None, **kwargs):
+        """Download model from hf hub and return the path."""
+        return hf_hub_download(
+            repo_type="model",
+            repo_id=repo_id,
+            filename=filename,
+        )
 
 def dynamic_load(root, model):
     module_path = f"{root.__name__}.{model}"

requirements.txt

@@ -2,7 +2,7 @@ e2cnn
 einops
 easydict
 gdown
-gradio==5.0.1
+# gradio==5.4.0
 h5py
 huggingface_hub
 imageio

third_party/ALIKE/assets/ALIKE_code.zip

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:891e8431c047e7aeed77c9e5f64ffeed262d92389d8ae6235dde0964a9048a08
-size 62774