vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/MatchAnything/imcui/hloc/matchers/sgmnet.py

@@ -0,0 +1,106 @@
+import sys
+from collections import OrderedDict, namedtuple
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, logger
+from ..utils.base_model import BaseModel
+
+sgmnet_path = Path(__file__).parent / "../../third_party/SGMNet"
+sys.path.append(str(sgmnet_path))
+
+from sgmnet import matcher as SGM_Model
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class SGMNet(BaseModel):
+    default_conf = {
+        "name": "SGM",
+        "model_name": "weights/sgm/root/model_best.pth",
+        "seed_top_k": [256, 256],
+        "seed_radius_coe": 0.01,
+        "net_channels": 128,
+        "layer_num": 9,
+        "head": 4,
+        "seedlayer": [0, 6],
+        "use_mc_seeding": True,
+        "use_score_encoding": False,
+        "conf_bar": [1.11, 0.1],
+        "sink_iter": [10, 100],
+        "detach_iter": 1000000,
+        "match_threshold": 0.2,
+    }
+    required_inputs = [
+        "image0",
+        "image1",
+    ]
+
+    # Initialize the line matcher
+    def _init(self, conf):
+        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(model_path, map_location="cpu")
+        # for ddp model
+        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
+            checkpoint["state_dict"] = new_stat_dict
+        self.net.load_state_dict(checkpoint["state_dict"])
+        logger.info("Load SGMNet model done.")
+
+    def _forward(self, data):
+        x1 = data["keypoints0"].squeeze()  # N x 2
+        x2 = data["keypoints1"].squeeze()
+        score1 = data["scores0"].reshape(-1, 1)  # N x 1
+        score2 = data["scores1"].reshape(-1, 1)
+        desc1 = data["descriptors0"].permute(0, 2, 1)  # 1 x N x 128
+        desc2 = data["descriptors1"].permute(0, 2, 1)
+        size1 = (
+            torch.tensor(data["image0"].shape[2:]).flip(0).to(x1.device)
+        )  # W x H -> x & y
+        size2 = torch.tensor(data["image1"].shape[2:]).flip(0).to(x2.device)  # W x H
+        norm_x1 = self.normalize_size(x1, size1)
+        norm_x2 = self.normalize_size(x2, size2)
+
+        x1 = torch.cat((norm_x1, score1), dim=-1)  # N x 3
+        x2 = torch.cat((norm_x2, score2), dim=-1)
+        input = {"x1": x1[None], "x2": x2[None], "desc1": desc1, "desc2": desc2}
+        input = {
+            k: v.to(device).float() if isinstance(v, torch.Tensor) else v
+            for k, v in input.items()
+        }
+        pred = self.net(input, test_mode=True)
+
+        p = pred["p"]  # shape: N * M
+        indices0 = self.match_p(p[0, :-1, :-1])
+        pred = {
+            "matches0": indices0.unsqueeze(0),
+            "matching_scores0": torch.zeros(indices0.size(0)).unsqueeze(0),
+        }
+        return pred
+
+    def match_p(self, p):
+        score, index = torch.topk(p, k=1, dim=-1)
+        _, index2 = torch.topk(p, k=1, dim=-2)
+        mask_th, index, index2 = (
+            score[:, 0] > self.conf["match_threshold"],
+            index[:, 0],
+            index2.squeeze(0),
+        )
+        mask_mc = index2[index] == torch.arange(len(p)).to(device)
+        mask = mask_th & mask_mc
+        indices0 = torch.where(mask, index, index.new_tensor(-1))
+        return indices0
+
+    def normalize_size(self, x, size, scale=1):
+        norm_fac = size.max()
+        return (x - size / 2 + 0.5) / (norm_fac * scale)

vismatch/third_party/MatchAnything/imcui/hloc/matchers/sold2.py

@@ -0,0 +1,144 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, logger
+from ..utils.base_model import BaseModel
+
+sold2_path = Path(__file__).parent / "../../third_party/SOLD2"
+sys.path.append(str(sold2_path))
+
+from sold2.model.line_matcher import LineMatcher
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class SOLD2(BaseModel):
+    default_conf = {
+        "model_name": "sold2_wireframe.tar",
+        "match_threshold": 0.2,
+        "checkpoint_dir": sold2_path / "pretrained",
+        "detect_thresh": 0.25,
+        "multiscale": False,
+        "valid_thresh": 1e-3,
+        "num_blocks": 20,
+        "overlap_ratio": 0.5,
+    }
+    required_inputs = [
+        "image0",
+        "image1",
+    ]
+
+    # Initialize the line matcher
+    def _init(self, conf):
+        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 = {
+            "model_cfg": {
+                "model_name": "lcnn_simple",
+                "model_architecture": "simple",
+                # Backbone related config
+                "backbone": "lcnn",
+                "backbone_cfg": {
+                    "input_channel": 1,  # Use RGB images or grayscale images.
+                    "depth": 4,
+                    "num_stacks": 2,
+                    "num_blocks": 1,
+                    "num_classes": 5,
+                },
+                # Junction decoder related config
+                "junction_decoder": "superpoint_decoder",
+                "junc_decoder_cfg": {},
+                # Heatmap decoder related config
+                "heatmap_decoder": "pixel_shuffle",
+                "heatmap_decoder_cfg": {},
+                # Descriptor decoder related config
+                "descriptor_decoder": "superpoint_descriptor",
+                "descriptor_decoder_cfg": {},
+                # Shared configurations
+                "grid_size": 8,
+                "keep_border_valid": True,
+                # Threshold of junction detection
+                "detection_thresh": 0.0153846,  # 1/65
+                "max_num_junctions": 300,
+                # Threshold of heatmap detection
+                "prob_thresh": 0.5,
+                # Weighting related parameters
+                "weighting_policy": mode,
+                # [Heatmap loss]
+                "w_heatmap": 0.0,
+                "w_heatmap_class": 1,
+                "heatmap_loss_func": "cross_entropy",
+                "heatmap_loss_cfg": {"policy": mode},
+                # [Heatmap consistency loss]
+                # [Junction loss]
+                "w_junc": 0.0,
+                "junction_loss_func": "superpoint",
+                "junction_loss_cfg": {"policy": mode},
+                # [Descriptor loss]
+                "w_desc": 0.0,
+                "descriptor_loss_func": "regular_sampling",
+                "descriptor_loss_cfg": {
+                    "dist_threshold": 8,
+                    "grid_size": 4,
+                    "margin": 1,
+                    "policy": mode,
+                },
+            },
+            "line_detector_cfg": {
+                "detect_thresh": 0.25,  # depending on your images, you might need to tune this parameter
+                "num_samples": 64,
+                "sampling_method": "local_max",
+                "inlier_thresh": 0.9,
+                "use_candidate_suppression": True,
+                "nms_dist_tolerance": 3.0,
+                "use_heatmap_refinement": True,
+                "heatmap_refine_cfg": {
+                    "mode": "local",
+                    "ratio": 0.2,
+                    "valid_thresh": 1e-3,
+                    "num_blocks": 20,
+                    "overlap_ratio": 0.5,
+                },
+            },
+            "multiscale": False,
+            "line_matcher_cfg": {
+                "cross_check": True,
+                "num_samples": 5,
+                "min_dist_pts": 8,
+                "top_k_candidates": 10,
+                "grid_size": 4,
+            },
+        }
+        self.net = LineMatcher(
+            match_config["model_cfg"],
+            model_path,
+            device,
+            match_config["line_detector_cfg"],
+            match_config["line_matcher_cfg"],
+            match_config["multiscale"],
+        )
+
+    def _forward(self, data):
+        img0 = data["image0"]
+        img1 = data["image1"]
+        pred = self.net([img0, img1])
+        line_seg1 = pred["line_segments"][0]
+        line_seg2 = pred["line_segments"][1]
+        matches = pred["matches"]
+
+        valid_matches = matches != -1
+        match_indices = matches[valid_matches]
+        matched_lines1 = line_seg1[valid_matches][:, :, ::-1]
+        matched_lines2 = line_seg2[match_indices][:, :, ::-1]
+
+        pred["raw_lines0"], pred["raw_lines1"] = line_seg1, line_seg2
+        pred["lines0"], pred["lines1"] = matched_lines1, matched_lines2
+        pred = {**pred, **data}
+        return pred

vismatch/third_party/MatchAnything/imcui/hloc/matchers/superglue.py

@@ -0,0 +1,33 @@
+import sys
+from pathlib import Path
+
+from ..utils.base_model import BaseModel
+
+sys.path.append(str(Path(__file__).parent / "../../third_party"))
+from SuperGluePretrainedNetwork.models.superglue import (  # noqa: E402
+    SuperGlue as SG,
+)
+
+
+class SuperGlue(BaseModel):
+    default_conf = {
+        "weights": "outdoor",
+        "sinkhorn_iterations": 100,
+        "match_threshold": 0.2,
+    }
+    required_inputs = [
+        "image0",
+        "keypoints0",
+        "scores0",
+        "descriptors0",
+        "image1",
+        "keypoints1",
+        "scores1",
+        "descriptors1",
+    ]
+
+    def _init(self, conf):
+        self.net = SG(conf)
+
+    def _forward(self, data):
+        return self.net(data)

vismatch/third_party/MatchAnything/imcui/hloc/matchers/topicfm.py

@@ -0,0 +1,60 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID
+
+from ..utils.base_model import BaseModel
+
+sys.path.append(str(Path(__file__).parent / "../../third_party"))
+from TopicFM.src import get_model_cfg
+from TopicFM.src.models.topic_fm import TopicFM as _TopicFM
+
+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,
+    }
+    required_inputs = ["image0", "image1"]
+
+    def _init(self, conf):
+        _conf = dict(get_model_cfg())
+        _conf["match_coarse"]["thr"] = conf["match_threshold"]
+        _conf["coarse"]["n_samples"] = conf["n_sampling_topics"]
+        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(model_path, map_location="cpu")
+        self.net.load_state_dict(ckpt_dict["state_dict"])
+
+    def _forward(self, data):
+        data_ = {
+            "image0": data["image0"],
+            "image1": data["image1"],
+        }
+        self.net(data_)
+        pred = {
+            "keypoints0": data_["mkpts0_f"],
+            "keypoints1": data_["mkpts1_f"],
+            "mconf": data_["mconf"],
+        }
+        scores = data_["mconf"]
+        top_k = self.conf["max_keypoints"]
+        if top_k is not None and len(scores) > top_k:
+            keep = torch.argsort(scores, descending=True)[:top_k]
+            scores = scores[keep]
+            pred["keypoints0"], pred["keypoints1"], pred["mconf"] = (
+                pred["keypoints0"][keep],
+                pred["keypoints1"][keep],
+                scores,
+            )
+        return pred

vismatch/third_party/MatchAnything/imcui/hloc/matchers/xfeat_dense.py

@@ -0,0 +1,54 @@
+import torch
+
+from .. import logger
+
+from ..utils.base_model import BaseModel
+
+
+class XFeatDense(BaseModel):
+    default_conf = {
+        "keypoint_threshold": 0.005,
+        "max_keypoints": 8000,
+    }
+    required_inputs = [
+        "image0",
+        "image1",
+    ]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "verlab/accelerated_features",
+            "XFeat",
+            pretrained=True,
+            top_k=self.conf["max_keypoints"],
+        )
+        logger.info("Load XFeat(dense) model done.")
+
+    def _forward(self, data):
+        # Compute coarse feats
+        out0 = self.net.detectAndComputeDense(
+            data["image0"], top_k=self.conf["max_keypoints"]
+        )
+        out1 = self.net.detectAndComputeDense(
+            data["image1"], top_k=self.conf["max_keypoints"]
+        )
+
+        # Match batches of pairs
+        idxs_list = self.net.batch_match(out0["descriptors"], out1["descriptors"])
+        B = len(data["image0"])
+
+        # Refine coarse matches
+        # this part is harder to batch, currently iterate
+        matches = []
+        for b in range(B):
+            matches.append(
+                self.net.refine_matches(out0, out1, matches=idxs_list, batch_idx=b)
+            )
+        # we use results from one batch
+        matches = matches[0]
+        pred = {
+            "keypoints0": matches[:, :2],
+            "keypoints1": matches[:, 2:],
+            "mconf": torch.ones_like(matches[:, 0]),
+        }
+        return pred

vismatch/third_party/MatchAnything/imcui/hloc/matchers/xfeat_lightglue.py

@@ -0,0 +1,48 @@
+import torch
+
+from .. import logger
+
+from ..utils.base_model import BaseModel
+
+
+class XFeatLightGlue(BaseModel):
+    default_conf = {
+        "keypoint_threshold": 0.005,
+        "max_keypoints": 8000,
+    }
+    required_inputs = [
+        "image0",
+        "image1",
+    ]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "verlab/accelerated_features",
+            "XFeat",
+            pretrained=True,
+            top_k=self.conf["max_keypoints"],
+        )
+        logger.info("Load XFeat(dense) model done.")
+
+    def _forward(self, data):
+        # we use results from one batch
+        im0 = data["image0"]
+        im1 = data["image1"]
+        # Compute coarse feats
+        out0 = self.net.detectAndCompute(im0, top_k=self.conf["max_keypoints"])[0]
+        out1 = self.net.detectAndCompute(im1, top_k=self.conf["max_keypoints"])[0]
+        out0.update({"image_size": (im0.shape[-1], im0.shape[-2])})  # W H
+        out1.update({"image_size": (im1.shape[-1], im1.shape[-2])})  # W H
+        pred = self.net.match_lighterglue(out0, out1)
+        if len(pred) == 3:
+            mkpts_0, mkpts_1, _ = pred
+        else:
+            mkpts_0, mkpts_1 = pred
+        mkpts_0 = torch.from_numpy(mkpts_0)  # n x 2
+        mkpts_1 = torch.from_numpy(mkpts_1)  # n x 2
+        pred = {
+            "keypoints0": mkpts_0,
+            "keypoints1": mkpts_1,
+            "mconf": torch.ones_like(mkpts_0[:, 0]),
+        }
+        return pred

vismatch/third_party/MatchAnything/imcui/hloc/matchers/xoftr.py

@@ -0,0 +1,90 @@
+import sys
+import warnings
+from pathlib import Path
+
+import torch
+
+from .. import DEVICE, MODEL_REPO_ID, logger
+
+tp_path = Path(__file__).parent / "../../third_party"
+sys.path.append(str(tp_path))
+
+from XoFTR.src.config.default import get_cfg_defaults
+from XoFTR.src.utils.misc import lower_config
+from XoFTR.src.xoftr import XoFTR as XoFTR_
+
+
+from ..utils.base_model import BaseModel
+
+
+class XoFTR(BaseModel):
+    default_conf = {
+        "model_name": "weights_xoftr_640.ckpt",
+        "match_threshold": 0.3,
+        "max_keypoints": -1,
+    }
+    required_inputs = ["image0", "image1"]
+
+    def _init(self, conf):
+        # Get default configurations
+        config_ = get_cfg_defaults(inference=True)
+        config_ = lower_config(config_)
+
+        # Coarse level threshold
+        config_["xoftr"]["match_coarse"]["thr"] = self.conf["match_threshold"]
+
+        # Fine level threshold
+        config_["xoftr"]["fine"]["thr"] = 0.1  # Default 0.1
+
+        # It is posseble to get denser matches
+        # If True, xoftr returns all fine-level matches for each fine-level window (at 1/2 resolution)
+        config_["xoftr"]["fine"]["denser"] = False  # Default False
+
+        # XoFTR model
+        matcher = XoFTR_(config=config_["xoftr"])
+
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+
+        # Load model
+        state_dict = torch.load(model_path, map_location="cpu")["state_dict"]
+        matcher.load_state_dict(state_dict, strict=True)
+        matcher = matcher.eval().to(DEVICE)
+        self.net = matcher
+        logger.info(f"Loaded XoFTR with weights {conf['model_name']}")
+
+    def _forward(self, data):
+        # For consistency with hloc pairs, we refine kpts in image0!
+        rename = {
+            "keypoints0": "keypoints1",
+            "keypoints1": "keypoints0",
+            "image0": "image1",
+            "image1": "image0",
+            "mask0": "mask1",
+            "mask1": "mask0",
+        }
+        data_ = {rename[k]: v for k, v in data.items()}
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            pred = self.net(data_)
+        pred = {
+            "keypoints0": data_["mkpts0_f"],
+            "keypoints1": data_["mkpts1_f"],
+        }
+        scores = data_["mconf_f"]
+
+        top_k = self.conf["max_keypoints"]
+        if top_k is not None and len(scores) > top_k:
+            keep = torch.argsort(scores, descending=True)[:top_k]
+            pred["keypoints0"], pred["keypoints1"] = (
+                pred["keypoints0"][keep],
+                pred["keypoints1"][keep],
+            )
+            scores = scores[keep]
+
+        # Switch back indices
+        pred = {(rename[k] if k in rename else k): v for k, v in pred.items()}
+        pred["scores"] = scores
+        return pred

vismatch/third_party/MatchAnything/imcui/hloc/pairs_from_covisibility.py

@@ -0,0 +1,60 @@
+import argparse
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+from tqdm import tqdm
+
+from . import logger
+from .utils.read_write_model import read_model
+
+
+def main(model, output, num_matched):
+    logger.info("Reading the COLMAP model...")
+    cameras, images, points3D = read_model(model)
+
+    logger.info("Extracting image pairs from covisibility info...")
+    pairs = []
+    for image_id, image in tqdm(images.items()):
+        matched = image.point3D_ids != -1
+        points3D_covis = image.point3D_ids[matched]
+
+        covis = defaultdict(int)
+        for point_id in points3D_covis:
+            for image_covis_id in points3D[point_id].image_ids:
+                if image_covis_id != image_id:
+                    covis[image_covis_id] += 1
+
+        if len(covis) == 0:
+            logger.info(f"Image {image_id} does not have any covisibility.")
+            continue
+
+        covis_ids = np.array(list(covis.keys()))
+        covis_num = np.array([covis[i] for i in covis_ids])
+
+        if len(covis_ids) <= num_matched:
+            top_covis_ids = covis_ids[np.argsort(-covis_num)]
+        else:
+            # get covisible image ids with top k number of common matches
+            ind_top = np.argpartition(covis_num, -num_matched)
+            ind_top = ind_top[-num_matched:]  # unsorted top k
+            ind_top = ind_top[np.argsort(-covis_num[ind_top])]
+            top_covis_ids = [covis_ids[i] for i in ind_top]
+            assert covis_num[ind_top[0]] == np.max(covis_num)
+
+        for i in top_covis_ids:
+            pair = (image.name, images[i].name)
+            pairs.append(pair)
+
+    logger.info(f"Found {len(pairs)} pairs.")
+    with open(output, "w") as f:
+        f.write("\n".join(" ".join([i, j]) for i, j in pairs))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True, type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument("--num_matched", required=True, type=int)
+    args = parser.parse_args()
+    main(**args.__dict__)

vismatch/third_party/MatchAnything/imcui/hloc/pairs_from_exhaustive.py

@@ -0,0 +1,64 @@
+import argparse
+import collections.abc as collections
+from pathlib import Path
+from typing import List, Optional, Union
+
+from . import logger
+from .utils.io import list_h5_names
+from .utils.parsers import parse_image_lists
+
+
+def main(
+    output: Path,
+    image_list: Optional[Union[Path, List[str]]] = None,
+    features: Optional[Path] = None,
+    ref_list: Optional[Union[Path, List[str]]] = None,
+    ref_features: Optional[Path] = None,
+):
+    if image_list is not None:
+        if isinstance(image_list, (str, Path)):
+            names_q = parse_image_lists(image_list)
+        elif isinstance(image_list, collections.Iterable):
+            names_q = list(image_list)
+        else:
+            raise ValueError(f"Unknown type for image list: {image_list}")
+    elif features is not None:
+        names_q = list_h5_names(features)
+    else:
+        raise ValueError("Provide either a list of images or a feature file.")
+
+    self_matching = False
+    if ref_list is not None:
+        if isinstance(ref_list, (str, Path)):
+            names_ref = parse_image_lists(ref_list)
+        elif isinstance(image_list, collections.Iterable):
+            names_ref = list(ref_list)
+        else:
+            raise ValueError(f"Unknown type for reference image list: {ref_list}")
+    elif ref_features is not None:
+        names_ref = list_h5_names(ref_features)
+    else:
+        self_matching = True
+        names_ref = names_q
+
+    pairs = []
+    for i, n1 in enumerate(names_q):
+        for j, n2 in enumerate(names_ref):
+            if self_matching and j <= i:
+                continue
+            pairs.append((n1, n2))
+
+    logger.info(f"Found {len(pairs)} pairs.")
+    with open(output, "w") as f:
+        f.write("\n".join(" ".join([i, j]) for i, j in pairs))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument("--image_list", type=Path)
+    parser.add_argument("--features", type=Path)
+    parser.add_argument("--ref_list", type=Path)
+    parser.add_argument("--ref_features", type=Path)
+    args = parser.parse_args()
+    main(**args.__dict__)