update: limit keypoints number

app.py

@@ -143,7 +143,7 @@ def run(config):
                         #     label="Matcher mode",
                         #     value="NN-mutual",
                         # )
-                    with gr.Accordion("RANSAC Setting", open=False):
+                    with gr.Accordion("RANSAC Setting", open=True):
                         with gr.Row(equal_height=False):
                             enable_ransac = gr.Checkbox(label="Enable RANSAC")
                             ransac_method = gr.Dropdown(
@@ -174,7 +174,7 @@ def run(config):
                             value=10000,
                         )
 
-                    with gr.Accordion("Geometry Setting", open=True):
+                    with gr.Accordion("Geometry Setting", open=False):
                         with gr.Row(equal_height=False):
                             # show_geom = gr.Checkbox(label="Show Geometry")
                             choice_estimate_geom = gr.Radio(
@@ -227,13 +227,13 @@ def run(config):
                     label="Keypoints Matching", type="numpy"
                 )
                 with gr.Accordion(
-                    "Open for More: Matches Statistics", open=False
+                    "Open for More: Matches Statistics", open=True
                 ):
                     matches_result_info = gr.JSON(label="Matches Statistics")
                     matcher_info = gr.JSON(label="Match info")
 
-                output_wrapped = gr.Image(label="Wrapped Pair", type="numpy")
                 with gr.Accordion("Open for More: Geometry info", open=False):
+                    output_wrapped = gr.Image(label="Wrapped Pair", type="numpy")
                     geometry_result = gr.JSON(label="Reconstructed Geometry")
 
             # callbacks

common/utils.py

@@ -53,7 +53,7 @@ def gen_examples():
     match_setting_threshold = 0.1
     match_setting_max_features = 2000
     detect_keypoints_threshold = 0.01
-    enable_ransac = False
+    enable_ransac = True
     ransac_method = "RANSAC"
     ransac_reproj_threshold = 8
     ransac_confidence = 0.999

hloc/__init__.py

@@ -4,7 +4,8 @@ from packaging import version
 __version__ = "1.3"
 
 formatter = logging.Formatter(
-    fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s", datefmt="%Y/%m/%d %H:%M:%S"
+    fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s",
+    datefmt="%Y/%m/%d %H:%M:%S",
 )
 handler = logging.StreamHandler()
 handler.setFormatter(formatter)

hloc/extract_features.py

@@ -85,17 +85,18 @@ confs = {
         "preprocessing": {
             "grayscale": False,
             "force_resize": True,
-            "resize_max": 1600,
+            "resize_max": 1024,
             "width": 640,
             "height": 480,
             "dfactor": 8,
         },
     },
     "d2net-ss": {
-        "output": "feats-d2net-ss",
+        "output": "feats-d2net-ss-n5000-r1600",
         "model": {
             "name": "d2net",
             "multiscale": False,
+            "max_keypoints": 5000,
         },
         "preprocessing": {
             "grayscale": False,
@@ -103,10 +104,11 @@ confs = {
         },
     },
     "d2net-ms": {
-        "output": "feats-d2net-ms",
+        "output": "feats-d2net-ms-n5000-r1600",
         "model": {
             "name": "d2net",
             "multiscale": True,
+            "max_keypoints": 5000,
         },
         "preprocessing": {
             "grayscale": False,
@@ -114,7 +116,7 @@ confs = {
         },
     },
     "rootsift": {
-        "output": "feats-sift",
+        "output": "feats-rootsift-n5000-r1600",
         "model": {
             "name": "dog",
             "max_keypoints": 5000,
@@ -129,7 +131,7 @@ confs = {
         },
     },
     "sift": {
-        "output": "feats-sift",
+        "output": "feats-sift-n5000-r1600",
         "model": {
             "name": "dog",
             "descriptor": "sift",
@@ -145,8 +147,12 @@ confs = {
         },
     },
     "sosnet": {
-        "output": "feats-sosnet",
-        "model": {"name": "dog", "descriptor": "sosnet"},
+        "output": "feats-sosnet-n5000-r1600",
+        "model": {
+            "name": "dog",
+            "descriptor": "sosnet",
+            "max_keypoints": 5000,
+        },
         "preprocessing": {
             "grayscale": True,
             "resize_max": 1600,
@@ -157,8 +163,12 @@ confs = {
         },
     },
     "hardnet": {
-        "output": "feats-hardnet",
-        "model": {"name": "dog", "descriptor": "hardnet"},
+        "output": "feats-hardnet-n5000-r1600",
+        "model": {
+            "name": "dog",
+            "descriptor": "hardnet",
+            "max_keypoints": 5000,
+        },
         "preprocessing": {
             "grayscale": True,
             "resize_max": 1600,
@@ -169,7 +179,7 @@ confs = {
         },
     },
     "disk": {
-        "output": "feats-disk",
+        "output": "feats-disk-n5000-r1600",
         "model": {
             "name": "disk",
             "max_keypoints": 5000,
@@ -180,7 +190,7 @@ confs = {
         },
     },
     "alike": {
-        "output": "feats-alike",
+        "output": "feats-alike-n5000-r1600",
         "model": {
             "name": "alike",
             "max_keypoints": 5000,
@@ -196,7 +206,7 @@ confs = {
         },
     },
     "lanet": {
-        "output": "feats-lanet",
+        "output": "feats-lanet-n5000-r1600",
         "model": {
             "name": "lanet",
             "keypoint_threshold": 0.1,
@@ -208,7 +218,7 @@ confs = {
         },
     },
     "darkfeat": {
-        "output": "feats-darkfeat-n5000-r1024",
+        "output": "feats-darkfeat-n5000-r1600",
         "model": {
             "name": "darkfeat",
             "max_keypoints": 5000,
@@ -225,7 +235,7 @@ confs = {
         },
     },
     "dedode": {
-        "output": "feats-dedode-n5000-r1024",
+        "output": "feats-dedode-n5000-r1600",
         "model": {
             "name": "dedode",
             "max_keypoints": 5000,
@@ -233,14 +243,14 @@ confs = {
         "preprocessing": {
             "grayscale": False,
             "force_resize": True,
-            "resize_max": 1024,
+            "resize_max": 1600,
             "width": 768,
             "height": 768,
             "dfactor": 8,
         },
     },
     "example": {
-        "output": "feats-example-n5000-r1024",
+        "output": "feats-example-n2000-r1024",
         "model": {
             "name": "example",
             "keypoint_threshold": 0.1,
@@ -323,13 +333,17 @@ class ImageDataset(torch.utils.data.Dataset):
             if isinstance(paths, (Path, str)):
                 self.names = parse_image_lists(paths)
             elif isinstance(paths, collections.Iterable):
-                self.names = [p.as_posix() if isinstance(p, Path) else p for p in paths]
+                self.names = [
+                    p.as_posix() if isinstance(p, Path) else p for p in paths
+                ]
             else:
                 raise ValueError(f"Unknown format for path argument {paths}.")
 
             for name in self.names:
                 if not (root / name).exists():
-                    raise ValueError(f"Image {name} does not exists in root: {root}.")
+                    raise ValueError(
+                        f"Image {name} does not exists in root: {root}."
+                    )
 
     def __getitem__(self, idx):
         name = self.names[idx]
@@ -397,7 +411,10 @@ def extract(model, image_0, conf):
 
         # assure that the size is divisible by dfactor
         size_new = tuple(
-            map(lambda x: int(x // conf.dfactor * conf.dfactor), image.shape[-2:])
+            map(
+                lambda x: int(x // conf.dfactor * conf.dfactor),
+                image.shape[-2:],
+            )
         )
         image = F.resize(image, size=size_new, antialias=True)
         input_ = image.to(device, non_blocking=True)[None]
@@ -435,7 +452,8 @@ def main(
     overwrite: bool = False,
 ) -> Path:
     logger.info(
-        "Extracting local features with configuration:" f"\n{pprint.pformat(conf)}"
+        "Extracting local features with configuration:"
+        f"\n{pprint.pformat(conf)}"
     )
 
     dataset = ImageDataset(image_dir, conf["preprocessing"], image_list)
@@ -443,7 +461,9 @@ def main(
         feature_path = Path(export_dir, conf["output"] + ".h5")
     feature_path.parent.mkdir(exist_ok=True, parents=True)
     skip_names = set(
-        list_h5_names(feature_path) if feature_path.exists() and not overwrite else ()
+        list_h5_names(feature_path)
+        if feature_path.exists() and not overwrite
+        else ()
     )
     dataset.names = [n for n in dataset.names if n not in skip_names]
     if len(dataset.names) == 0:
@@ -507,7 +527,10 @@ if __name__ == "__main__":
     parser.add_argument("--image_dir", type=Path, required=True)
     parser.add_argument("--export_dir", type=Path, required=True)
     parser.add_argument(
-        "--conf", type=str, default="superpoint_aachen", choices=list(confs.keys())
+        "--conf",
+        type=str,
+        default="superpoint_aachen",
+        choices=list(confs.keys()),
     )
     parser.add_argument("--as_half", action="store_true")
     parser.add_argument("--image_list", type=Path)

hloc/extractors/d2net.py

@@ -17,6 +17,7 @@ class D2Net(BaseModel):
         "checkpoint_dir": d2net_path / "models",
         "use_relu": True,
         "multiscale": False,
+        "max_keypoints": 1024,
     }
     required_inputs = ["image"]
 
@@ -50,6 +51,11 @@ class D2Net(BaseModel):
             )
         keypoints = keypoints[:, [1, 0]]  # (x, y) and remove the scale
 
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        keypoints = keypoints[idxs, :2]
+        descriptors = descriptors[idxs]
+        scores = scores[idxs]
+
         return {
             "keypoints": torch.from_numpy(keypoints)[None],
             "scores": torch.from_numpy(scores)[None],

hloc/extractors/darkfeat.py

@@ -32,7 +32,9 @@ class DarkFeat(BaseModel):
             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}`.")
+            logger.info(
+                f"Downloading the DarkFeat model with `{cmd_wo_proxy}`."
+            )
             try:
                 subprocess.run(cmd_wo_proxy, check=True)
             except subprocess.CalledProcessError as e:
@@ -50,6 +52,10 @@ class DarkFeat(BaseModel):
         keypoints = pred["keypoints"]
         descriptors = pred["descriptors"]
         scores = pred["scores"]
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        keypoints = keypoints[idxs, :2]
+        descriptors = descriptors[:, idxs]
+        scores = scores[idxs]
         return {
             "keypoints": keypoints[None],  # 1 x N x 2
             "scores": scores[None],  # 1 x N

hloc/extractors/dedode.py

@@ -36,7 +36,9 @@ class DeDoDe(BaseModel):
 
     # Initialize the line matcher
     def _init(self, conf):
-        model_detector_path = dedode_path / "pretrained" / conf["model_detector_name"]
+        model_detector_path = (
+            dedode_path / "pretrained" / conf["model_detector_name"]
+        )
         model_descriptor_path = (
             dedode_path / "pretrained" / conf["model_descriptor_name"]
         )
@@ -56,17 +58,24 @@ class DeDoDe(BaseModel):
             model_descriptor_path.parent.mkdir(exist_ok=True)
             link = self.weight_urls[conf["model_descriptor_name"]]
             cmd = ["wget", link, "-O", str(model_descriptor_path)]
-            logger.info(f"Downloading the DeDoDe descriptor model with `{cmd}`.")
+            logger.info(
+                f"Downloading the DeDoDe descriptor model with `{cmd}`."
+            )
             subprocess.run(cmd, check=True)
 
         logger.info(f"Loading DeDoDe model...")
 
         # load the model
         weights_detector = torch.load(model_detector_path, map_location="cpu")
-        weights_descriptor = torch.load(model_descriptor_path, map_location="cpu")
-        self.detector = dedode_detector_L(weights=weights_detector, device=device)
-        self.descriptor = dedode_descriptor_B(weights=weights_descriptor, device=device)
-
+        weights_descriptor = torch.load(
+            model_descriptor_path, map_location="cpu"
+        )
+        self.detector = dedode_detector_L(
+            weights=weights_detector, device=device
+        )
+        self.descriptor = dedode_descriptor_B(
+            weights=weights_descriptor, device=device
+        )
         logger.info(f"Load DeDoDe model done.")
 
     def _forward(self, data):
@@ -91,9 +100,9 @@ class DeDoDe(BaseModel):
 
         # step 2: describe keypoints
         # dim: 1 x N x 256
-        description_A = self.descriptor.describe_keypoints(batch_A, keypoints_A)[
-            "descriptions"
-        ]
+        description_A = self.descriptor.describe_keypoints(
+            batch_A, keypoints_A
+        )["descriptions"]
         keypoints_A = to_pixel_coords(keypoints_A, H_A, W_A)
 
         return {

hloc/extractors/dir.py

@@ -8,7 +8,9 @@ import gdown
 
 from ..utils.base_model import BaseModel
 
-sys.path.append(str(Path(__file__).parent / "../../third_party/deep-image-retrieval"))
+sys.path.append(
+    str(Path(__file__).parent / "../../third_party/deep-image-retrieval")
+)
 os.environ["DB_ROOT"] = ""  # required by dirtorch
 
 from dirtorch.utils import common  # noqa: E402
@@ -40,7 +42,9 @@ class DIR(BaseModel):
     }
 
     def _init(self, conf):
-        checkpoint = Path(torch.hub.get_dir(), "dirtorch", conf["model_name"] + ".pt")
+        checkpoint = Path(
+            torch.hub.get_dir(), "dirtorch", conf["model_name"] + ".pt"
+        )
         if not checkpoint.exists():
             checkpoint.parent.mkdir(exist_ok=True, parents=True)
             link = self.dir_models[conf["model_name"]]

hloc/extractors/dog.py

@@ -1,5 +1,8 @@
 import kornia
-from kornia.feature.laf import laf_from_center_scale_ori, extract_patches_from_pyramid
+from kornia.feature.laf import (
+    laf_from_center_scale_ori,
+    extract_patches_from_pyramid,
+)
 import numpy as np
 import torch
 import pycolmap

hloc/extractors/example.py

@@ -26,7 +26,6 @@ class Example(BaseModel):
     required_inputs = ["image"]
 
     def _init(self, conf):
-
         # set checkpoints paths if needed
         model_path = example_path / "checkpoints" / f'{conf["model_name"]}'
         if not model_path.exists():

hloc/extractors/fire.py

@@ -34,7 +34,6 @@ class FIRe(BaseModel):
     }
 
     def _init(self, conf):
-
         assert conf["model_name"] in self.fire_models.keys()
         # Config paths
         model_path = fire_path / "model" / conf["model_name"]
@@ -64,7 +63,6 @@ class FIRe(BaseModel):
         self.scales = conf["scales"]
 
     def _forward(self, data):
-
         image = self.norm_rgb(data["image"])
 
         # Feature extraction.

hloc/extractors/fire_local.py

@@ -41,7 +41,6 @@ class FIRe(BaseModel):
     }
 
     def _init(self, conf):
-
         assert conf["model_name"] in self.fire_models.keys()
 
         # Config paths
@@ -75,7 +74,6 @@ class FIRe(BaseModel):
         self.features_num = conf["features_num"]
 
     def _forward(self, data):
-
         image = self.norm_rgb(data["image"])
 
         local_desc = self.net.forward_local(

hloc/extractors/lanet.py

@@ -21,7 +21,9 @@ class LANet(BaseModel):
     required_inputs = ["image"]
 
     def _init(self, conf):
-        model_path = lanet_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        model_path = (
+            lanet_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        )
         if not model_path.exists():
             print(f"No model found at {model_path}")
         self.net = PointModel(is_test=True)
@@ -34,16 +36,16 @@ class LANet(BaseModel):
         _, _, Hc, Wc = descriptors.shape
 
         # Scores & Descriptors
-        kpts_score = (
-            torch.cat([keypoints, scores], dim=1).view(3, -1).t()
-        )
-        descriptors = (
-            descriptors.view(256, Hc, Wc).view(256, -1).t()
-        )
+        kpts_score = torch.cat([keypoints, scores], dim=1).view(3, -1).t()
+        descriptors = descriptors.view(256, Hc, Wc).view(256, -1).t()
 
         # Filter based on confidence threshold
-        descriptors = descriptors[kpts_score[:, 0] > self.conf["keypoint_threshold"], :]
-        kpts_score = kpts_score[kpts_score[:, 0] > self.conf["keypoint_threshold"], :]
+        descriptors = descriptors[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
+        kpts_score = kpts_score[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
         keypoints = kpts_score[:, 1:]
         scores = kpts_score[:, 0]
 

hloc/extractors/netvlad.py

@@ -18,7 +18,9 @@ EPS = 1e-6
 class NetVLADLayer(nn.Module):
     def __init__(self, input_dim=512, K=64, score_bias=False, intranorm=True):
         super().__init__()
-        self.score_proj = nn.Conv1d(input_dim, K, kernel_size=1, bias=score_bias)
+        self.score_proj = nn.Conv1d(
+            input_dim, K, kernel_size=1, bias=score_bias
+        )
         centers = nn.parameter.Parameter(torch.empty([input_dim, K]))
         nn.init.xavier_uniform_(centers)
         self.register_parameter("centers", centers)
@@ -54,7 +56,9 @@ class NetVLAD(BaseModel):
         assert conf["model_name"] in self.dir_models.keys()
 
         # Download the checkpoint.
-        checkpoint = Path(torch.hub.get_dir(), "netvlad", conf["model_name"] + ".mat")
+        checkpoint = Path(
+            torch.hub.get_dir(), "netvlad", conf["model_name"] + ".mat"
+        )
         if not checkpoint.exists():
             checkpoint.parent.mkdir(exist_ok=True, parents=True)
             link = self.dir_models[conf["model_name"]]
@@ -77,7 +81,9 @@ class NetVLAD(BaseModel):
         mat = loadmat(checkpoint, struct_as_record=False, squeeze_me=True)
 
         # CNN weights.
-        for layer, mat_layer in zip(self.backbone.children(), mat["net"].layers):
+        for layer, mat_layer in zip(
+            self.backbone.children(), mat["net"].layers
+        ):
             if isinstance(layer, nn.Conv2d):
                 w = mat_layer.weights[0]  # Shape: S x S x IN x OUT
                 b = mat_layer.weights[1]  # Shape: OUT

hloc/extractors/rekd.py

@@ -20,7 +20,9 @@ class REKD(BaseModel):
     required_inputs = ["image"]
 
     def _init(self, conf):
-        model_path = rekd_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        model_path = (
+            rekd_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth'
+        )
         if not model_path.exists():
             print(f"No model found at {model_path}")
         self.net = REKD_(is_test=True)
@@ -34,15 +36,29 @@ class REKD(BaseModel):
 
         # Scores & Descriptors
         kpts_score = (
-            torch.cat([keypoints, scores], dim=1).view(3, -1).t().cpu().detach().numpy()
+            torch.cat([keypoints, scores], dim=1)
+            .view(3, -1)
+            .t()
+            .cpu()
+            .detach()
+            .numpy()
         )
         descriptors = (
-            descriptors.view(256, Hc, Wc).view(256, -1).t().cpu().detach().numpy()
+            descriptors.view(256, Hc, Wc)
+            .view(256, -1)
+            .t()
+            .cpu()
+            .detach()
+            .numpy()
         )
 
         # Filter based on confidence threshold
-        descriptors = descriptors[kpts_score[:, 0] > self.conf["keypoint_threshold"], :]
-        kpts_score = kpts_score[kpts_score[:, 0] > self.conf["keypoint_threshold"], :]
+        descriptors = descriptors[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
+        kpts_score = kpts_score[
+            kpts_score[:, 0] > self.conf["keypoint_threshold"], :
+        ]
         keypoints = kpts_score[:, 1:]
         scores = kpts_score[:, 0]
 

hloc/extractors/superpoint.py

@@ -16,7 +16,10 @@ def sample_descriptors_fix_sampling(keypoints, descriptors, s: int = 8):
     keypoints = (keypoints + 0.5) / (keypoints.new_tensor([w, h]) * s)
     keypoints = keypoints * 2 - 1  # normalize to (-1, 1)
     descriptors = torch.nn.functional.grid_sample(
-        descriptors, keypoints.view(b, 1, -1, 2), mode="bilinear", align_corners=False
+        descriptors,
+        keypoints.view(b, 1, -1, 2),
+        mode="bilinear",
+        align_corners=False,
     )
     descriptors = torch.nn.functional.normalize(
         descriptors.reshape(b, c, -1), p=2, dim=1

hloc/match_dense.py

@@ -224,7 +224,10 @@ def match(model, path_0, path_1, conf):
         image = torch.from_numpy(image / 255.0).float()
         # assure that the size is divisible by dfactor
         size_new = tuple(
-            map(lambda x: int(x // conf.dfactor * conf.dfactor), image.shape[-2:])
+            map(
+                lambda x: int(x // conf.dfactor * conf.dfactor),
+                image.shape[-2:],
+            )
         )
         image = F.resize(image, size=size_new, antialias=True)
         scale = np.array(size) / np.array(size_new)[::-1]
@@ -291,7 +294,10 @@ def match_images(model, image_0, image_1, conf, device="cpu"):
 
         # assure that the size is divisible by dfactor
         size_new = tuple(
-            map(lambda x: int(x // conf.dfactor * conf.dfactor), image.shape[-2:])
+            map(
+                lambda x: int(x // conf.dfactor * conf.dfactor),
+                image.shape[-2:],
+            )
         )
         image = F.resize(image, size=size_new)
         scale = np.array(size) / np.array(size_new)[::-1]
@@ -348,7 +354,7 @@ def match_images(model, image_0, image_1, conf, device="cpu"):
         if "mconf" in pred.keys():
             ret["mconf"] = pred["mconf"].cpu().numpy()
         else:
-            ret["mconf"] = np.ones_like(kpts0.cpu().numpy()[:,0])
+            ret["mconf"] = np.ones_like(kpts0.cpu().numpy()[:, 0])
     if "lines0" in pred.keys() and "lines1" in pred.keys():
         if "keypoints0" in pred.keys() and "keypoints1" in pred.keys():
             kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
@@ -357,7 +363,10 @@ def match_images(model, image_0, image_1, conf, device="cpu"):
             kpts0_origin = kpts0_origin.cpu().numpy()
             kpts1_origin = kpts1_origin.cpu().numpy()
         else:
-            kpts0_origin, kpts1_origin = None, None  # np.zeros([0]), np.zeros([0])
+            kpts0_origin, kpts1_origin = (
+                None,
+                None,
+            )  # np.zeros([0]), np.zeros([0])
         lines0, lines1 = pred["lines0"], pred["lines1"]
         lines0_raw, lines1_raw = pred["raw_lines0"], pred["raw_lines1"]
 

hloc/match_features.py

@@ -151,7 +151,8 @@ class WorkQueue:
     def __init__(self, work_fn, num_threads=1):
         self.queue = Queue(num_threads)
         self.threads = [
-            Thread(target=self.thread_fn, args=(work_fn,)) for _ in range(num_threads)
+            Thread(target=self.thread_fn, args=(work_fn,))
+            for _ in range(num_threads)
         ]
         for thread in self.threads:
             thread.start()
@@ -220,21 +221,24 @@ def main(
     features_ref: Optional[Path] = None,
     overwrite: bool = False,
 ) -> Path:
-
     if isinstance(features, Path) or Path(features).exists():
         features_q = features
         if matches is None:
             raise ValueError(
-                "Either provide both features and matches as Path" " or both as names."
+                "Either provide both features and matches as Path"
+                " or both as names."
             )
     else:
         if export_dir is None:
             raise ValueError(
-                "Provide an export_dir if features is not" f" a file path: {features}."
+                "Provide an export_dir if features is not"
+                f" a file path: {features}."
             )
         features_q = Path(export_dir, features + ".h5")
         if matches is None:
-            matches = Path(export_dir, f'{features}_{conf["output"]}_{pairs.stem}.h5')
+            matches = Path(
+                export_dir, f'{features}_{conf["output"]}_{pairs.stem}.h5'
+            )
 
     if features_ref is None:
         features_ref = features_q
@@ -276,7 +280,8 @@ def match_from_paths(
     overwrite: bool = False,
 ) -> Path:
     logger.info(
-        "Matching local features with configuration:" f"\n{pprint.pformat(conf)}"
+        "Matching local features with configuration:"
+        f"\n{pprint.pformat(conf)}"
     )
 
     if not feature_path_q.exists():
@@ -330,12 +335,11 @@ def match_images(model, feat0, feat1):
         desc0 = desc0.unsqueeze(0)
     if len(desc1.shape) == 2:
         desc1 = desc1.unsqueeze(0)
-
     if isinstance(feat0["keypoints"], list):
         feat0["keypoints"] = feat0["keypoints"][0][None]
     if isinstance(feat1["keypoints"], list):
         feat1["keypoints"] = feat1["keypoints"][0][None]
-    
+
     pred = model(
         {
             "image0": feat0["image"],
@@ -386,7 +390,9 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--pairs", type=Path, required=True)
     parser.add_argument("--export_dir", type=Path)
-    parser.add_argument("--features", type=str, default="feats-superpoint-n4096-r1024")
+    parser.add_argument(
+        "--features", type=str, default="feats-superpoint-n4096-r1024"
+    )
     parser.add_argument("--matches", type=Path)
     parser.add_argument(
         "--conf", type=str, default="superglue", choices=list(confs.keys())

hloc/matchers/aspanformer.py

@@ -21,6 +21,7 @@ class ASpanFormer(BaseModel):
     default_conf = {
         "weights": "outdoor",
         "match_threshold": 0.2,
+        "sinkhorn_iterations": 20,
         "config_path": aspanformer_path / "configs/aspan/outdoor/aspan_test.py",
         "model_name": "weights_aspanformer.tar",
     }
@@ -31,24 +32,39 @@ class ASpanFormer(BaseModel):
     }
 
     def _init(self, conf):
-        model_path = aspanformer_path / "weights" / Path(conf["weights"] + ".ckpt")
+        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 = [
+                    "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}`.")
+                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 the Aspanformer model with `{cmd}`.")
+                    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.")
+                        logger.error(
+                            f"Failed to download the Aspanformer model."
+                        )
                         raise e
 
             do_system(f"cd {str(aspanformer_path)} & tar -xvf {str(tar_path)}")
@@ -58,9 +74,16 @@ class ASpanFormer(BaseModel):
         config = get_cfg_defaults()
         config.merge_from_file(conf["config_path"])
         _config = lower_config(config)
+
+        # update: match threshold
+        _config["aspan"]["match_coarse"]["thr"] = conf["match_threshold"]
+        _config["aspan"]["match_coarse"]["skh_iters"] = conf["sinkhorn_iterations"]
+
         self.net = _ASpanFormer(config=_config["aspan"])
         weight_path = model_path
-        state_dict = torch.load(str(weight_path), map_location="cpu")["state_dict"]
+        state_dict = torch.load(str(weight_path), map_location="cpu")[
+            "state_dict"
+        ]
         self.net.load_state_dict(state_dict, strict=False)
 
     def _forward(self, data):

hloc/matchers/dkm.py

@@ -55,7 +55,9 @@ class DKMv3(BaseModel):
 
         warp, certainty = self.net.match(img0, img1, device=device)
         matches, certainty = self.net.sample(warp, certainty)
-        kpts1, kpts2 = self.net.to_pixel_coordinates(matches, H_A, W_A, H_B, W_B)
+        kpts1, kpts2 = self.net.to_pixel_coordinates(
+            matches, H_A, W_A, H_B, W_B
+        )
         pred = {}
         pred["keypoints0"], pred["keypoints1"] = kpts1, kpts2
         return pred

hloc/matchers/dual_softmax.py

@@ -3,6 +3,7 @@ import torch
 from ..utils.base_model import BaseModel
 import numpy as np
 
+
 # borrow from dedode
 def dual_softmax_matcher(
     desc_A: tuple["B", "C", "N"],
@@ -17,7 +18,9 @@ def dual_softmax_matcher(
     if normalize:
         desc_A = desc_A / desc_A.norm(dim=1, keepdim=True)
         desc_B = desc_B / desc_B.norm(dim=1, keepdim=True)
-    sim = torch.einsum("b c n, b c m -> b n m", desc_A, desc_B) * inv_temperature
+    sim = (
+        torch.einsum("b c n, b c m -> b n m", desc_A, desc_B) * inv_temperature
+    )
     P = sim.softmax(dim=-2) * sim.softmax(dim=-1)
     mask = torch.nonzero(
         (P == P.max(dim=-1, keepdim=True).values)
@@ -47,9 +50,14 @@ class DualSoftMax(BaseModel):
         pass
 
     def _forward(self, data):
-        if data["descriptors0"].size(-1) == 0 or data["descriptors1"].size(-1) == 0:
+        if (
+            data["descriptors0"].size(-1) == 0
+            or data["descriptors1"].size(-1) == 0
+        ):
             matches0 = torch.full(
-                data["descriptors0"].shape[:2], -1, device=data["descriptors0"].device
+                data["descriptors0"].shape[:2],
+                -1,
+                device=data["descriptors0"].device,
             )
             return {
                 "matches0": matches0,

hloc/matchers/gluestick.py

@@ -33,9 +33,12 @@ class GlueStick(BaseModel):
     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"]
+        model_path = (
+            gluestick_path / "resources" / "weights" / conf["model_name"]
+        )
 
         # Download the model.
         if not model_path.exists():

hloc/matchers/nearest_neighbor.py

@@ -36,24 +36,36 @@ class NearestNeighbor(BaseModel):
         pass
 
     def _forward(self, data):
-        if data["descriptors0"].size(-1) == 0 or data["descriptors1"].size(-1) == 0:
+        if (
+            data["descriptors0"].size(-1) == 0
+            or data["descriptors1"].size(-1) == 0
+        ):
             matches0 = torch.full(
-                data["descriptors0"].shape[:2], -1, device=data["descriptors0"].device
+                data["descriptors0"].shape[:2],
+                -1,
+                device=data["descriptors0"].device,
             )
             return {
                 "matches0": matches0,
                 "matching_scores0": torch.zeros_like(matches0),
             }
         ratio_threshold = self.conf["ratio_threshold"]
-        if data["descriptors0"].size(-1) == 1 or data["descriptors1"].size(-1) == 1:
+        if (
+            data["descriptors0"].size(-1) == 1
+            or data["descriptors1"].size(-1) == 1
+        ):
             ratio_threshold = None
-        sim = torch.einsum("bdn,bdm->bnm", data["descriptors0"], data["descriptors1"])
+        sim = torch.einsum(
+            "bdn,bdm->bnm", data["descriptors0"], data["descriptors1"]
+        )
         matches0, scores0 = find_nn(
             sim, ratio_threshold, self.conf["distance_threshold"]
         )
         if self.conf["do_mutual_check"]:
             matches1, scores1 = find_nn(
-                sim.transpose(1, 2), ratio_threshold, self.conf["distance_threshold"]
+                sim.transpose(1, 2),
+                ratio_threshold,
+                self.conf["distance_threshold"],
             )
             matches0 = mutual_check(matches0, matches1)
         return {

hloc/matchers/roma.py

@@ -84,7 +84,9 @@ class Roma(BaseModel):
         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)
+        kpts1, kpts2 = self.net.to_pixel_coordinates(
+            matches, H_A, W_A, H_B, W_B
+        )
         pred = {}
         pred["keypoints0"], pred["keypoints1"] = kpts1, kpts2
         pred["mconf"] = certainty

hloc/matchers/sgmnet.py

@@ -52,9 +52,18 @@ class SGMNet(BaseModel):
         # Download the model.
         if not sgmnet_weights.exists():
             if not tar_path.exists():
-                cmd = ["gdown", link, "-O", str(tar_path), "--proxy", self.proxy]
+                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}`.")
+                logger.info(
+                    f"Downloading the SGMNet model with `{cmd_wo_proxy}`."
+                )
                 try:
                     subprocess.run(cmd_wo_proxy, check=True)
                 except subprocess.CalledProcessError as e:
@@ -73,7 +82,10 @@ class SGMNet(BaseModel):
         self.net = SGM_Model(config)
         checkpoint = torch.load(sgmnet_weights, map_location="cpu")
         # for ddp model
-        if list(checkpoint["state_dict"].items())[0][0].split(".")[0] == "module":
+        if (
+            list(checkpoint["state_dict"].items())[0][0].split(".")[0]
+            == "module"
+        ):
             new_stat_dict = OrderedDict()
             for key, value in checkpoint["state_dict"].items():
                 new_stat_dict[key[7:]] = value

hloc/matchers/sold2.py

@@ -35,6 +35,7 @@ class SOLD2(BaseModel):
         "image0",
         "image1",
     ]
+
     # Initialize the line matcher
     def _init(self, conf):
         checkpoint_path = conf["checkpoint_dir"] / conf["weights"]

hloc/pipelines/4Seasons/localize.py

@@ -67,7 +67,9 @@ delete_unused_images(seq_images, timestamps)
 generate_query_lists(timestamps, seq_dir, query_list)
 
 # Generate the localization pairs from the given reference frames.
-generate_localization_pairs(sequence, reloc, num_loc_pairs, ref_pairs, loc_pairs)
+generate_localization_pairs(
+    sequence, reloc, num_loc_pairs, ref_pairs, loc_pairs
+)
 
 # Extract, match, amd localize.
 ffile = extract_features.main(fconf, seq_images, output_dir)

hloc/pipelines/4Seasons/utils.py

@@ -48,7 +48,11 @@ def camera_from_calibration_file(id_, path):
     model_name = "PINHOLE"
     params = [float(i) for i in [fx, fy, cx, cy]]
     camera = Camera(
-        id=id_, model=model_name, width=int(width), height=int(height), params=params
+        id=id_,
+        model=model_name,
+        width=int(width),
+        height=int(height),
+        params=params,
     )
     return camera
 
@@ -149,7 +153,9 @@ def generate_localization_pairs(sequence, reloc, num, ref_pairs, out_path):
     """
     if "test" in sequence:
         # hard pairs will be overwritten by easy ones if available
-        relocs = [str(reloc).replace("*", d) for d in ["hard", "moderate", "easy"]]
+        relocs = [
+            str(reloc).replace("*", d) for d in ["hard", "moderate", "easy"]
+        ]
     else:
         relocs = [reloc]
     query_to_ref_ts = {}
@@ -207,8 +213,12 @@ def evaluate_submission(submission_dir, relocs, ths=[0.1, 0.2, 0.5]):
     """Compute the relocalization recall from predicted and ground truth poses."""
     for reloc in relocs.parent.glob(relocs.name):
         poses_gt = parse_relocalization(reloc, has_poses=True)
-        poses_pred = parse_relocalization(submission_dir / reloc.name, has_poses=True)
-        poses_pred = {(ref_ts, q_ts): (R, t) for ref_ts, q_ts, R, t in poses_pred}
+        poses_pred = parse_relocalization(
+            submission_dir / reloc.name, has_poses=True
+        )
+        poses_pred = {
+            (ref_ts, q_ts): (R, t) for ref_ts, q_ts, R, t in poses_pred
+        }
 
         error = []
         for ref_ts, q_ts, R_gt, t_gt in poses_gt:

hloc/pipelines/7Scenes/create_gt_sfm.py

@@ -28,7 +28,9 @@ def interpolate_depth(depth, kp):
 
     # To maximize the number of points that have depth:
     # do bilinear interpolation first and then nearest for the remaining points
-    interp_lin = grid_sample(depth, kp, align_corners=True, mode="bilinear")[0, :, 0]
+    interp_lin = grid_sample(depth, kp, align_corners=True, mode="bilinear")[
+        0, :, 0
+    ]
     interp_nn = torch.nn.functional.grid_sample(
         depth, kp, align_corners=True, mode="nearest"
     )[0, :, 0]
@@ -127,7 +129,15 @@ if __name__ == "__main__":
     dataset = Path("datasets/7scenes")
     outputs = Path("outputs/7Scenes")
 
-    SCENES = ["chess", "fire", "heads", "office", "pumpkin", "redkitchen", "stairs"]
+    SCENES = [
+        "chess",
+        "fire",
+        "heads",
+        "office",
+        "pumpkin",
+        "redkitchen",
+        "stairs",
+    ]
     for scene in SCENES:
         sfm_path = outputs / scene / "sfm_superpoint+superglue"
         depth_path = dataset / f"depth/7scenes_{scene}/train/depth"

hloc/pipelines/7Scenes/pipeline.py

@@ -45,7 +45,9 @@ def run_scene(
     create_reference_sfm(gt_dir, ref_sfm_sift, test_list)
     create_query_list_with_intrinsics(gt_dir, query_list, test_list)
 
-    features = extract_features.main(feature_conf, images, outputs, as_half=True)
+    features = extract_features.main(
+        feature_conf, images, outputs, as_half=True
+    )
 
     sfm_pairs = outputs / f"pairs-db-covis{num_covis}.txt"
     pairs_from_covisibility.main(ref_sfm_sift, sfm_pairs, num_matched=num_covis)
@@ -112,7 +114,9 @@ if __name__ == "__main__":
         results = (
             args.outputs
             / scene
-            / "results_{}.txt".format("dense" if args.use_dense_depth else "sparse")
+            / "results_{}.txt".format(
+                "dense" if args.use_dense_depth else "sparse"
+            )
         )
         if args.overwrite or not results.exists():
             run_scene(

hloc/pipelines/Aachen/pipeline.py

@@ -40,14 +40,18 @@ images = dataset / "images/images_upright/"
 
 outputs = args.outputs  # where everything will be saved
 sift_sfm = outputs / "sfm_sift"  # from which we extract the reference poses
-reference_sfm = outputs / "sfm_superpoint+superglue"  # the SfM model we will build
+reference_sfm = (
+    outputs / "sfm_superpoint+superglue"
+)  # the SfM model we will build
 sfm_pairs = (
     outputs / f"pairs-db-covis{args.num_covis}.txt"
 )  # top-k most covisible in SIFT model
 loc_pairs = (
     outputs / f"pairs-query-netvlad{args.num_loc}.txt"
 )  # top-k retrieved by NetVLAD
-results = outputs / f"Aachen_hloc_superpoint+superglue_netvlad{args.num_loc}.txt"
+results = (
+    outputs / f"Aachen_hloc_superpoint+superglue_netvlad{args.num_loc}.txt"
+)
 
 # list the standard configurations available
 print(f"Configs for feature extractors:\n{pformat(extract_features.confs)}")
@@ -71,7 +75,9 @@ sfm_matches = match_features.main(
     matcher_conf, sfm_pairs, feature_conf["output"], outputs
 )
 
-triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches)
+triangulation.main(
+    reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches
+)
 
 global_descriptors = extract_features.main(retrieval_conf, images, outputs)
 pairs_from_retrieval.main(

hloc/pipelines/Aachen_v1_1/pipeline.py

@@ -39,14 +39,18 @@ images = dataset / "images/images_upright/"
 sift_sfm = dataset / "3D-models/aachen_v_1_1"
 
 outputs = args.outputs  # where everything will be saved
-reference_sfm = outputs / "sfm_superpoint+superglue"  # the SfM model we will build
+reference_sfm = (
+    outputs / "sfm_superpoint+superglue"
+)  # the SfM model we will build
 sfm_pairs = (
     outputs / f"pairs-db-covis{args.num_covis}.txt"
 )  # top-k most covisible in SIFT model
 loc_pairs = (
     outputs / f"pairs-query-netvlad{args.num_loc}.txt"
 )  # top-k retrieved by NetVLAD
-results = outputs / f"Aachen-v1.1_hloc_superpoint+superglue_netvlad{args.num_loc}.txt"
+results = (
+    outputs / f"Aachen-v1.1_hloc_superpoint+superglue_netvlad{args.num_loc}.txt"
+)
 
 # list the standard configurations available
 print(f"Configs for feature extractors:\n{pformat(extract_features.confs)}")
@@ -64,7 +68,9 @@ sfm_matches = match_features.main(
     matcher_conf, sfm_pairs, feature_conf["output"], outputs
 )
 
-triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches)
+triangulation.main(
+    reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches
+)
 
 global_descriptors = extract_features.main(retrieval_conf, images, outputs)
 pairs_from_retrieval.main(

hloc/pipelines/Aachen_v1_1/pipeline_loftr.py

@@ -61,7 +61,9 @@ features, sfm_matches = match_dense.main(
     matcher_conf, sfm_pairs, images, outputs, max_kps=8192, overwrite=False
 )
 
-triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches)
+triangulation.main(
+    reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches
+)
 
 global_descriptors = extract_features.main(retrieval_conf, images, outputs)
 pairs_from_retrieval.main(

hloc/pipelines/CMU/pipeline.py

@@ -46,20 +46,34 @@ def run_slice(slice_, root, outputs, num_covis, num_loc):
     matcher_conf = match_features.confs["superglue"]
 
     pairs_from_covisibility.main(sift_sfm, sfm_pairs, num_matched=num_covis)
-    features = extract_features.main(feature_conf, ref_images, outputs, as_half=True)
+    features = extract_features.main(
+        feature_conf, ref_images, outputs, as_half=True
+    )
     sfm_matches = match_features.main(
         matcher_conf, sfm_pairs, feature_conf["output"], outputs
     )
-    triangulation.main(ref_sfm, sift_sfm, ref_images, sfm_pairs, features, sfm_matches)
+    triangulation.main(
+        ref_sfm, sift_sfm, ref_images, sfm_pairs, features, sfm_matches
+    )
 
     generate_query_list(root, query_list, slice_)
-    global_descriptors = extract_features.main(retrieval_conf, ref_images, outputs)
-    global_descriptors = extract_features.main(retrieval_conf, query_images, outputs)
+    global_descriptors = extract_features.main(
+        retrieval_conf, ref_images, outputs
+    )
+    global_descriptors = extract_features.main(
+        retrieval_conf, query_images, outputs
+    )
     pairs_from_retrieval.main(
-        global_descriptors, loc_pairs, num_loc, query_list=query_list, db_model=ref_sfm
+        global_descriptors,
+        loc_pairs,
+        num_loc,
+        query_list=query_list,
+        db_model=ref_sfm,
     )
 
-    features = extract_features.main(feature_conf, query_images, outputs, as_half=True)
+    features = extract_features.main(
+        feature_conf, query_images, outputs, as_half=True
+    )
     loc_matches = match_features.main(
         matcher_conf, loc_pairs, feature_conf["output"], outputs
     )
@@ -122,5 +136,9 @@ if __name__ == "__main__":
     for slice_ in slices:
         logger.info("Working on slice %s.", slice_)
         run_slice(
-            f"slice{slice_}", args.dataset, args.outputs, args.num_covis, args.num_loc
+            f"slice{slice_}",
+            args.dataset,
+            args.outputs,
+            args.num_covis,
+            args.num_loc,
         )

hloc/pipelines/Cambridge/pipeline.py

@@ -5,7 +5,13 @@ from .utils import create_query_list_with_intrinsics, scale_sfm_images, evaluate
 from ... import extract_features, match_features, pairs_from_covisibility
 from ... import triangulation, localize_sfm, pairs_from_retrieval, logger
 
-SCENES = ["KingsCollege", "OldHospital", "ShopFacade", "StMarysChurch", "GreatCourt"]
+SCENES = [
+    "KingsCollege",
+    "OldHospital",
+    "ShopFacade",
+    "StMarysChurch",
+    "GreatCourt",
+]
 
 
 def run_scene(images, gt_dir, outputs, results, num_covis, num_loc):
@@ -35,7 +41,11 @@ def run_scene(images, gt_dir, outputs, results, num_covis, num_loc):
     retrieval_conf = extract_features.confs["netvlad"]
 
     create_query_list_with_intrinsics(
-        gt_dir / "empty_all", query_list, test_list, ext=".txt", image_dir=images
+        gt_dir / "empty_all",
+        query_list,
+        test_list,
+        ext=".txt",
+        image_dir=images,
     )
     with open(test_list, "r") as f:
         query_seqs = {q.split("/")[0] for q in f.read().rstrip().split("\n")}
@@ -49,7 +59,9 @@ def run_scene(images, gt_dir, outputs, results, num_covis, num_loc):
         query_prefix=query_seqs,
     )
 
-    features = extract_features.main(feature_conf, images, outputs, as_half=True)
+    features = extract_features.main(
+        feature_conf, images, outputs, as_half=True
+    )
     pairs_from_covisibility.main(ref_sfm_sift, sfm_pairs, num_matched=num_covis)
     sfm_matches = match_features.main(
         matcher_conf, sfm_pairs, feature_conf["output"], outputs

hloc/pipelines/Cambridge/utils.py

@@ -42,7 +42,9 @@ def scale_sfm_images(full_model, scaled_model, image_dir):
         sy = h / camera.height
         assert sx == sy, (sx, sy)
         scaled_cameras[cam_id] = camera._replace(
-            width=w, height=h, params=camera.params * np.array([sx, sx, sy, 1.0])
+            width=w,
+            height=h,
+            params=camera.params * np.array([sx, sx, sy, 1.0]),
         )
 
     write_model(scaled_cameras, images, points3D, scaled_model)

hloc/pipelines/RobotCar/colmap_from_nvm.py

@@ -16,11 +16,14 @@ from ...utils.read_write_model import write_model
 logger = logging.getLogger(__name__)
 
 
-def read_nvm_model(nvm_path, database_path, image_ids, camera_ids, skip_points=False):
-
+def read_nvm_model(
+    nvm_path, database_path, image_ids, camera_ids, skip_points=False
+):
     # Extract the intrinsics from the db file instead of the NVM model
     db = sqlite3.connect(str(database_path))
-    ret = db.execute("SELECT camera_id, model, width, height, params FROM cameras;")
+    ret = db.execute(
+        "SELECT camera_id, model, width, height, params FROM cameras;"
+    )
     cameras = {}
     for camera_id, camera_model, width, height, params in ret:
         params = np.fromstring(params, dtype=np.double).reshape(-1)

hloc/pipelines/RobotCar/pipeline.py

@@ -79,7 +79,9 @@ sift_sfm = outputs / "sfm_sift"
 reference_sfm = outputs / "sfm_superpoint+superglue"
 sfm_pairs = outputs / f"pairs-db-covis{args.num_covis}.txt"
 loc_pairs = outputs / f"pairs-query-netvlad{args.num_loc}.txt"
-results = outputs / f"RobotCar_hloc_superpoint+superglue_netvlad{args.num_loc}.txt"
+results = (
+    outputs / f"RobotCar_hloc_superpoint+superglue_netvlad{args.num_loc}.txt"
+)
 
 # pick one of the configurations for extraction and matching
 retrieval_conf = extract_features.confs["netvlad"]
@@ -103,7 +105,9 @@ sfm_matches = match_features.main(
     matcher_conf, sfm_pairs, feature_conf["output"], outputs
 )
 
-triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches)
+triangulation.main(
+    reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches
+)
 
 global_descriptors = extract_features.main(retrieval_conf, images, outputs)
 # TODO: do per location and per camera

hloc/utils/database.py

@@ -100,7 +100,9 @@ CREATE_MATCHES_TABLE = """CREATE TABLE IF NOT EXISTS matches (
     cols INTEGER NOT NULL,
     data BLOB)"""
 
-CREATE_NAME_INDEX = "CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)"
+CREATE_NAME_INDEX = (
+    "CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)"
+)
 
 CREATE_ALL = "; ".join(
     [
@@ -150,20 +152,34 @@ class COLMAPDatabase(sqlite3.Connection):
         super(COLMAPDatabase, self).__init__(*args, **kwargs)
 
         self.create_tables = lambda: self.executescript(CREATE_ALL)
-        self.create_cameras_table = lambda: self.executescript(CREATE_CAMERAS_TABLE)
+        self.create_cameras_table = lambda: self.executescript(
+            CREATE_CAMERAS_TABLE
+        )
         self.create_descriptors_table = lambda: self.executescript(
             CREATE_DESCRIPTORS_TABLE
         )
-        self.create_images_table = lambda: self.executescript(CREATE_IMAGES_TABLE)
+        self.create_images_table = lambda: self.executescript(
+            CREATE_IMAGES_TABLE
+        )
         self.create_two_view_geometries_table = lambda: self.executescript(
             CREATE_TWO_VIEW_GEOMETRIES_TABLE
         )
-        self.create_keypoints_table = lambda: self.executescript(CREATE_KEYPOINTS_TABLE)
-        self.create_matches_table = lambda: self.executescript(CREATE_MATCHES_TABLE)
+        self.create_keypoints_table = lambda: self.executescript(
+            CREATE_KEYPOINTS_TABLE
+        )
+        self.create_matches_table = lambda: self.executescript(
+            CREATE_MATCHES_TABLE
+        )
         self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)
 
     def add_camera(
-        self, model, width, height, params, prior_focal_length=False, camera_id=None
+        self,
+        model,
+        width,
+        height,
+        params,
+        prior_focal_length=False,
+        camera_id=None,
     ):
         params = np.asarray(params, np.float64)
         cursor = self.execute(
@@ -298,7 +314,12 @@ def example_usage():
 
     # Create dummy cameras.
 
-    model1, width1, height1, params1 = 0, 1024, 768, np.array((1024.0, 512.0, 384.0))
+    model1, width1, height1, params1 = (
+        0,
+        1024,
+        768,
+        np.array((1024.0, 512.0, 384.0)),
+    )
     model2, width2, height2, params2 = (
         2,
         1024,

hloc/utils/geometry.py

@@ -16,12 +16,12 @@ def compute_epipolar_errors(qvec_r2t, tvec_r2t, p2d_r, p2d_t):
     E = vector_to_cross_product_matrix(T_r2t[:3, -1]) @ T_r2t[:3, :3]
     l2d_r2t = (E @ to_homogeneous(p2d_r).T).T
     l2d_t2r = (E.T @ to_homogeneous(p2d_t).T).T
-    errors_r = np.abs(np.sum(to_homogeneous(p2d_r) * l2d_t2r, axis=1)) / np.linalg.norm(
-        l2d_t2r[:, :2], axis=1
-    )
-    errors_t = np.abs(np.sum(to_homogeneous(p2d_t) * l2d_r2t, axis=1)) / np.linalg.norm(
-        l2d_r2t[:, :2], axis=1
-    )
+    errors_r = np.abs(
+        np.sum(to_homogeneous(p2d_r) * l2d_t2r, axis=1)
+    ) / np.linalg.norm(l2d_t2r[:, :2], axis=1)
+    errors_t = np.abs(
+        np.sum(to_homogeneous(p2d_t) * l2d_r2t, axis=1)
+    ) / np.linalg.norm(l2d_r2t[:, :2], axis=1)
     return E, errors_r, errors_t
 
 

hloc/utils/read_write_model.py

@@ -42,7 +42,9 @@ logger = logging.getLogger(__name__)
 CameraModel = collections.namedtuple(
     "CameraModel", ["model_id", "model_name", "num_params"]
 )
-Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"])
+Camera = collections.namedtuple(
+    "Camera", ["id", "model", "width", "height", "params"]
+)
 BaseImage = collections.namedtuple(
     "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]
 )
@@ -126,7 +128,11 @@ def read_cameras_text(path):
                 height = int(elems[3])
                 params = np.array(tuple(map(float, elems[4:])))
                 cameras[camera_id] = Camera(
-                    id=camera_id, model=model, width=width, height=height, params=params
+                    id=camera_id,
+                    model=model,
+                    width=width,
+                    height=height,
+                    params=params,
                 )
     return cameras
 
@@ -151,7 +157,9 @@ def read_cameras_binary(path_to_model_file):
             height = camera_properties[3]
             num_params = CAMERA_MODEL_IDS[model_id].num_params
             params = read_next_bytes(
-                fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params
+                fid,
+                num_bytes=8 * num_params,
+                format_char_sequence="d" * num_params,
             )
             cameras[camera_id] = Camera(
                 id=camera_id,
@@ -222,7 +230,10 @@ def read_images_text(path):
                 image_name = elems[9]
                 elems = fid.readline().split()
                 xys = np.column_stack(
-                    [tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))]
+                    [
+                        tuple(map(float, elems[0::3])),
+                        tuple(map(float, elems[1::3])),
+                    ]
                 )
                 point3D_ids = np.array(tuple(map(int, elems[2::3])))
                 images[image_id] = Image(
@@ -259,16 +270,19 @@ def read_images_binary(path_to_model_file):
             while current_char != b"\x00":  # look for the ASCII 0 entry
                 image_name += current_char.decode("utf-8")
                 current_char = read_next_bytes(fid, 1, "c")[0]
-            num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[
-                0
-            ]
+            num_points2D = read_next_bytes(
+                fid, num_bytes=8, format_char_sequence="Q"
+            )[0]
             x_y_id_s = read_next_bytes(
                 fid,
                 num_bytes=24 * num_points2D,
                 format_char_sequence="ddq" * num_points2D,
             )
             xys = np.column_stack(
-                [tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))]
+                [
+                    tuple(map(float, x_y_id_s[0::3])),
+                    tuple(map(float, x_y_id_s[1::3])),
+                ]
             )
             point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
             images[image_id] = Image(
@@ -307,7 +321,13 @@ def write_images_text(images, path):
     with open(path, "w") as fid:
         fid.write(HEADER)
         for _, img in images.items():
-            image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name]
+            image_header = [
+                img.id,
+                *img.qvec,
+                *img.tvec,
+                img.camera_id,
+                img.name,
+            ]
             first_line = " ".join(map(str, image_header))
             fid.write(first_line + "\n")
 
@@ -387,9 +407,9 @@ def read_points3D_binary(path_to_model_file):
             xyz = np.array(binary_point_line_properties[1:4])
             rgb = np.array(binary_point_line_properties[4:7])
             error = np.array(binary_point_line_properties[7])
-            track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[
-                0
-            ]
+            track_length = read_next_bytes(
+                fid, num_bytes=8, format_char_sequence="Q"
+            )[0]
             track_elems = read_next_bytes(
                 fid,
                 num_bytes=8 * track_length,
@@ -478,8 +498,12 @@ def read_model(path, ext=""):
             ext = ".txt"
         else:
             try:
-                cameras, images, points3D = read_model(os.path.join(path, "model/"))
-                logger.warning("This SfM file structure was deprecated in hloc v1.1")
+                cameras, images, points3D = read_model(
+                    os.path.join(path, "model/")
+                )
+                logger.warning(
+                    "This SfM file structure was deprecated in hloc v1.1"
+                )
                 return cameras, images, points3D
             except FileNotFoundError:
                 raise FileNotFoundError(
@@ -571,7 +595,9 @@ def main():
     )
     args = parser.parse_args()
 
-    cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format)
+    cameras, images, points3D = read_model(
+        path=args.input_model, ext=args.input_format
+    )
 
     print("num_cameras:", len(cameras))
     print("num_images:", len(images))
@@ -579,7 +605,11 @@ def main():
 
     if args.output_model is not None:
         write_model(
-            cameras, images, points3D, path=args.output_model, ext=args.output_format
+            cameras,
+            images,
+            points3D,
+            path=args.output_model,
+            ext=args.output_format,
         )
 
 

hloc/utils/viz.py

@@ -19,7 +19,9 @@ def cm_RdGn(x):
     return np.clip(c, 0, 1)
 
 
-def plot_images(imgs, titles=None, cmaps="gray", dpi=100, pad=0.5, adaptive=True):
+def plot_images(
+    imgs, titles=None, cmaps="gray", dpi=100, pad=0.5, adaptive=True
+):
     """Plot a set of images horizontally.
     Args:
         imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W).
@@ -129,7 +131,13 @@ def add_text(
 ):
     ax = plt.gcf().axes[idx]
     t = ax.text(
-        *pos, text, fontsize=fs, ha=ha, va=va, color=color, transform=ax.transAxes
+        *pos,
+        text,
+        fontsize=fs,
+        ha=ha,
+        va=va,
+        color=color,
+        transform=ax.transAxes
     )
     if lcolor is not None:
         t.set_path_effects(

hloc/utils/viz_3d.py

@@ -46,7 +46,9 @@ def init_figure(height: int = 800) -> go.Figure:
             dragmode="orbit",
         ),
         margin=dict(l=0, r=0, b=0, t=0, pad=0),
-        legend=dict(orientation="h", yanchor="top", y=0.99, xanchor="left", x=0.1),
+        legend=dict(
+            orientation="h", yanchor="top", y=0.99, xanchor="left", x=0.1
+        ),
     )
     return fig
 
@@ -68,7 +70,9 @@ def plot_points(
         mode="markers",
         name=name,
         legendgroup=name,
-        marker=dict(size=ps, color=color, line_width=0.0, colorscale=colorscale),
+        marker=dict(
+            size=ps, color=color, line_width=0.0, colorscale=colorscale
+        ),
     )
     fig.add_trace(tr)
 
@@ -162,7 +166,9 @@ def plot_camera_colmap(
     )
 
 
-def plot_cameras(fig: go.Figure, reconstruction: pycolmap.Reconstruction, **kwargs):
+def plot_cameras(
+    fig: go.Figure, reconstruction: pycolmap.Reconstruction, **kwargs
+):
     """Plot a camera as a cone with camera frustum."""
     for image_id, image in reconstruction.images.items():
         plot_camera_colmap(