vismatch 1.1.1__py3-none-any.whl

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

@@ -0,0 +1,459 @@
+import argparse
+import pprint
+from functools import partial
+from pathlib import Path
+from queue import Queue
+from threading import Thread
+from typing import Dict, List, Optional, Tuple, Union
+
+import h5py
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from . import logger, matchers
+from .utils.base_model import dynamic_load
+from .utils.parsers import names_to_pair, names_to_pair_old, parse_retrieval
+
+"""
+A set of standard configurations that can be directly selected from the command
+line using their name. Each is a dictionary with the following entries:
+    - output: the name of the match file that will be generated.
+    - model: the model configuration, as passed to a feature matcher.
+"""
+confs = {
+    "superglue": {
+        "output": "matches-superglue",
+        "model": {
+            "name": "superglue",
+            "weights": "outdoor",
+            "sinkhorn_iterations": 50,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "superglue-fast": {
+        "output": "matches-superglue-it5",
+        "model": {
+            "name": "superglue",
+            "weights": "outdoor",
+            "sinkhorn_iterations": 5,
+            "match_threshold": 0.2,
+        },
+    },
+    "superpoint-lightglue": {
+        "output": "matches-lightglue",
+        "model": {
+            "name": "lightglue",
+            "match_threshold": 0.2,
+            "width_confidence": 0.99,  # for point pruning
+            "depth_confidence": 0.95,  # for early stopping,
+            "features": "superpoint",
+            "model_name": "superpoint_lightglue.pth",
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "disk-lightglue": {
+        "output": "matches-disk-lightglue",
+        "model": {
+            "name": "lightglue",
+            "match_threshold": 0.2,
+            "width_confidence": 0.99,  # for point pruning
+            "depth_confidence": 0.95,  # for early stopping,
+            "features": "disk",
+            "model_name": "disk_lightglue.pth",
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "aliked-lightglue": {
+        "output": "matches-aliked-lightglue",
+        "model": {
+            "name": "lightglue",
+            "match_threshold": 0.2,
+            "width_confidence": 0.99,  # for point pruning
+            "depth_confidence": 0.95,  # for early stopping,
+            "features": "aliked",
+            "model_name": "aliked_lightglue.pth",
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "sift-lightglue": {
+        "output": "matches-sift-lightglue",
+        "model": {
+            "name": "lightglue",
+            "match_threshold": 0.2,
+            "width_confidence": 0.99,  # for point pruning
+            "depth_confidence": 0.95,  # for early stopping,
+            "features": "sift",
+            "add_scale_ori": True,
+            "model_name": "sift_lightglue.pth",
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "sgmnet": {
+        "output": "matches-sgmnet",
+        "model": {
+            "name": "sgmnet",
+            "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,
+        },
+        "preprocessing": {
+            "grayscale": True,
+            "resize_max": 1024,
+            "dfactor": 8,
+            "force_resize": False,
+        },
+    },
+    "NN-superpoint": {
+        "output": "matches-NN-mutual-dist.7",
+        "model": {
+            "name": "nearest_neighbor",
+            "do_mutual_check": True,
+            "distance_threshold": 0.7,
+            "match_threshold": 0.2,
+        },
+    },
+    "NN-ratio": {
+        "output": "matches-NN-mutual-ratio.8",
+        "model": {
+            "name": "nearest_neighbor",
+            "do_mutual_check": True,
+            "ratio_threshold": 0.8,
+            "match_threshold": 0.2,
+        },
+    },
+    "NN-mutual": {
+        "output": "matches-NN-mutual",
+        "model": {
+            "name": "nearest_neighbor",
+            "do_mutual_check": True,
+            "match_threshold": 0.2,
+        },
+    },
+    "Dual-Softmax": {
+        "output": "matches-Dual-Softmax",
+        "model": {
+            "name": "dual_softmax",
+            "match_threshold": 0.01,
+            "inv_temperature": 20,
+        },
+    },
+    "adalam": {
+        "output": "matches-adalam",
+        "model": {
+            "name": "adalam",
+            "match_threshold": 0.2,
+        },
+    },
+    "imp": {
+        "output": "matches-imp",
+        "model": {
+            "name": "imp",
+            "match_threshold": 0.2,
+        },
+    },
+}
+
+
+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)
+        ]
+        for thread in self.threads:
+            thread.start()
+
+    def join(self):
+        for thread in self.threads:
+            self.queue.put(None)
+        for thread in self.threads:
+            thread.join()
+
+    def thread_fn(self, work_fn):
+        item = self.queue.get()
+        while item is not None:
+            work_fn(item)
+            item = self.queue.get()
+
+    def put(self, data):
+        self.queue.put(data)
+
+
+class FeaturePairsDataset(torch.utils.data.Dataset):
+    def __init__(self, pairs, feature_path_q, feature_path_r):
+        self.pairs = pairs
+        self.feature_path_q = feature_path_q
+        self.feature_path_r = feature_path_r
+
+    def __getitem__(self, idx):
+        name0, name1 = self.pairs[idx]
+        data = {}
+        with h5py.File(self.feature_path_q, "r") as fd:
+            grp = fd[name0]
+            for k, v in grp.items():
+                data[k + "0"] = torch.from_numpy(v.__array__()).float()
+            # some matchers might expect an image but only use its size
+            data["image0"] = torch.empty((1,) + tuple(grp["image_size"])[::-1])
+        with h5py.File(self.feature_path_r, "r") as fd:
+            grp = fd[name1]
+            for k, v in grp.items():
+                data[k + "1"] = torch.from_numpy(v.__array__()).float()
+            data["image1"] = torch.empty((1,) + tuple(grp["image_size"])[::-1])
+        return data
+
+    def __len__(self):
+        return len(self.pairs)
+
+
+def writer_fn(inp, match_path):
+    pair, pred = inp
+    with h5py.File(str(match_path), "a", libver="latest") as fd:
+        if pair in fd:
+            del fd[pair]
+        grp = fd.create_group(pair)
+        matches = pred["matches0"][0].cpu().short().numpy()
+        grp.create_dataset("matches0", data=matches)
+        if "matching_scores0" in pred:
+            scores = pred["matching_scores0"][0].cpu().half().numpy()
+            grp.create_dataset("matching_scores0", data=scores)
+
+
+def main(
+    conf: Dict,
+    pairs: Path,
+    features: Union[Path, str],
+    export_dir: Optional[Path] = None,
+    matches: Optional[Path] = None,
+    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."
+            )
+    else:
+        if export_dir is None:
+            raise ValueError(
+                "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')
+
+    if features_ref is None:
+        features_ref = features_q
+    match_from_paths(conf, pairs, matches, features_q, features_ref, overwrite)
+
+    return matches
+
+
+def find_unique_new_pairs(pairs_all: List[Tuple[str]], match_path: Path = None):
+    """Avoid to recompute duplicates to save time."""
+    pairs = set()
+    for i, j in pairs_all:
+        if (j, i) not in pairs:
+            pairs.add((i, j))
+    pairs = list(pairs)
+    if match_path is not None and match_path.exists():
+        with h5py.File(str(match_path), "r", libver="latest") as fd:
+            pairs_filtered = []
+            for i, j in pairs:
+                if (
+                    names_to_pair(i, j) in fd
+                    or names_to_pair(j, i) in fd
+                    or names_to_pair_old(i, j) in fd
+                    or names_to_pair_old(j, i) in fd
+                ):
+                    continue
+                pairs_filtered.append((i, j))
+        return pairs_filtered
+    return pairs
+
+
+@torch.no_grad()
+def match_from_paths(
+    conf: Dict,
+    pairs_path: Path,
+    match_path: Path,
+    feature_path_q: Path,
+    feature_path_ref: Path,
+    overwrite: bool = False,
+) -> Path:
+    logger.info(
+        "Matching local features with configuration:" f"\n{pprint.pformat(conf)}"
+    )
+
+    if not feature_path_q.exists():
+        raise FileNotFoundError(f"Query feature file {feature_path_q}.")
+    if not feature_path_ref.exists():
+        raise FileNotFoundError(f"Reference feature file {feature_path_ref}.")
+    match_path.parent.mkdir(exist_ok=True, parents=True)
+
+    assert pairs_path.exists(), pairs_path
+    pairs = parse_retrieval(pairs_path)
+    pairs = [(q, r) for q, rs in pairs.items() for r in rs]
+    pairs = find_unique_new_pairs(pairs, None if overwrite else match_path)
+    if len(pairs) == 0:
+        logger.info("Skipping the matching.")
+        return
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    Model = dynamic_load(matchers, conf["model"]["name"])
+    model = Model(conf["model"]).eval().to(device)
+
+    dataset = FeaturePairsDataset(pairs, feature_path_q, feature_path_ref)
+    loader = torch.utils.data.DataLoader(
+        dataset, num_workers=5, batch_size=1, shuffle=False, pin_memory=True
+    )
+    writer_queue = WorkQueue(partial(writer_fn, match_path=match_path), 5)
+
+    for idx, data in enumerate(tqdm(loader, smoothing=0.1)):
+        data = {
+            k: v if k.startswith("image") else v.to(device, non_blocking=True)
+            for k, v in data.items()
+        }
+        pred = model(data)
+        pair = names_to_pair(*pairs[idx])
+        writer_queue.put((pair, pred))
+    writer_queue.join()
+    logger.info("Finished exporting matches.")
+
+
+def scale_keypoints(kpts, scale):
+    if (
+        isinstance(scale, (list, tuple, np.ndarray))
+        and len(scale) == 2
+        and np.any(scale != np.array([1.0, 1.0]))
+    ):
+        if isinstance(kpts, torch.Tensor):
+            kpts[:, 0] *= scale[0]  # scale x-dimension
+            kpts[:, 1] *= scale[1]  # scale y-dimension
+        elif isinstance(kpts, np.ndarray):
+            kpts[:, 0] *= scale[0]  # scale x-dimension
+            kpts[:, 1] *= scale[1]  # scale y-dimension
+    return kpts
+
+
+@torch.no_grad()
+def match_images(model, feat0, feat1):
+    # forward pass to match keypoints
+    desc0 = feat0["descriptors"][0]
+    desc1 = feat1["descriptors"][0]
+    if len(desc0.shape) == 2:
+        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]
+    input_dict = {
+        "image0": feat0["image"],
+        "keypoints0": feat0["keypoints"],
+        "scores0": feat0["scores"][0].unsqueeze(0),
+        "descriptors0": desc0,
+        "image1": feat1["image"],
+        "keypoints1": feat1["keypoints"],
+        "scores1": feat1["scores"][0].unsqueeze(0),
+        "descriptors1": desc1,
+    }
+    if "scales" in feat0:
+        input_dict = {**input_dict, "scales0": feat0["scales"]}
+    if "scales" in feat1:
+        input_dict = {**input_dict, "scales1": feat1["scales"]}
+    if "oris" in feat0:
+        input_dict = {**input_dict, "oris0": feat0["oris"]}
+    if "oris" in feat1:
+        input_dict = {**input_dict, "oris1": feat1["oris"]}
+    pred = model(input_dict)
+    pred = {
+        k: v.cpu().detach()[0] if isinstance(v, torch.Tensor) else v
+        for k, v in pred.items()
+    }
+    kpts0, kpts1 = (
+        feat0["keypoints"][0].cpu().numpy(),
+        feat1["keypoints"][0].cpu().numpy(),
+    )
+    matches, confid = pred["matches0"], pred["matching_scores0"]
+    # Keep the matching keypoints.
+    valid = matches > -1
+    mkpts0 = kpts0[valid]
+    mkpts1 = kpts1[matches[valid]]
+    mconfid = confid[valid]
+    # rescale the keypoints to their original size
+    s0 = feat0["original_size"] / feat0["size"]
+    s1 = feat1["original_size"] / feat1["size"]
+    kpts0_origin = scale_keypoints(torch.from_numpy(kpts0 + 0.5), s0) - 0.5
+    kpts1_origin = scale_keypoints(torch.from_numpy(kpts1 + 0.5), s1) - 0.5
+
+    mkpts0_origin = scale_keypoints(torch.from_numpy(mkpts0 + 0.5), s0) - 0.5
+    mkpts1_origin = scale_keypoints(torch.from_numpy(mkpts1 + 0.5), s1) - 0.5
+
+    ret = {
+        "image0_orig": feat0["image_orig"],
+        "image1_orig": feat1["image_orig"],
+        "keypoints0": kpts0,
+        "keypoints1": kpts1,
+        "keypoints0_orig": kpts0_origin.numpy(),
+        "keypoints1_orig": kpts1_origin.numpy(),
+        "mkeypoints0": mkpts0,
+        "mkeypoints1": mkpts1,
+        "mkeypoints0_orig": mkpts0_origin.numpy(),
+        "mkeypoints1_orig": mkpts1_origin.numpy(),
+        "mconf": mconfid.numpy(),
+    }
+    del feat0, feat1, desc0, desc1, kpts0, kpts1, kpts0_origin, kpts1_origin
+    torch.cuda.empty_cache()
+
+    return ret
+
+
+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("--matches", type=Path)
+    parser.add_argument(
+        "--conf", type=str, default="superglue", choices=list(confs.keys())
+    )
+    args = parser.parse_args()
+    main(confs[args.conf], args.pairs, args.features, args.export_dir)

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

@@ -0,0 +1,3 @@
+def get_matcher(matcher):
+    mod = __import__(f"{__name__}.{matcher}", fromlist=[""])
+    return getattr(mod, "Model")

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

@@ -0,0 +1,68 @@
+import torch
+from kornia.feature.adalam import AdalamFilter
+from kornia.utils.helpers import get_cuda_device_if_available
+
+from ..utils.base_model import BaseModel
+
+
+class AdaLAM(BaseModel):
+    # See https://kornia.readthedocs.io/en/latest/_modules/kornia/feature/adalam/adalam.html.
+    default_conf = {
+        "area_ratio": 100,
+        "search_expansion": 4,
+        "ransac_iters": 128,
+        "min_inliers": 6,
+        "min_confidence": 200,
+        "orientation_difference_threshold": 30,
+        "scale_rate_threshold": 1.5,
+        "detected_scale_rate_threshold": 5,
+        "refit": True,
+        "force_seed_mnn": True,
+        "device": get_cuda_device_if_available(),
+    }
+    required_inputs = [
+        "image0",
+        "image1",
+        "descriptors0",
+        "descriptors1",
+        "keypoints0",
+        "keypoints1",
+        "scales0",
+        "scales1",
+        "oris0",
+        "oris1",
+    ]
+
+    def _init(self, conf):
+        self.adalam = AdalamFilter(conf)
+
+    def _forward(self, data):
+        assert data["keypoints0"].size(0) == 1
+        if data["keypoints0"].size(1) < 2 or data["keypoints1"].size(1) < 2:
+            matches = torch.zeros(
+                (0, 2), dtype=torch.int64, device=data["keypoints0"].device
+            )
+        else:
+            matches = self.adalam.match_and_filter(
+                data["keypoints0"][0],
+                data["keypoints1"][0],
+                data["descriptors0"][0].T,
+                data["descriptors1"][0].T,
+                data["image0"].shape[2:],
+                data["image1"].shape[2:],
+                data["oris0"][0],
+                data["oris1"][0],
+                data["scales0"][0],
+                data["scales1"][0],
+            )
+        matches_new = torch.full(
+            (data["keypoints0"].size(1),),
+            -1,
+            dtype=torch.int64,
+            device=data["keypoints0"].device,
+        )
+        matches_new[matches[:, 0]] = matches[:, 1]
+        return {
+            "matches0": matches_new.unsqueeze(0),
+            "matching_scores0": torch.zeros(matches_new.size(0)).unsqueeze(0),
+        }

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

@@ -0,0 +1,66 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, logger
+from ..utils.base_model import BaseModel
+
+sys.path.append(str(Path(__file__).parent / "../../third_party"))
+from ASpanFormer.src.ASpanFormer.aspanformer import ASpanFormer as _ASpanFormer
+from ASpanFormer.src.config.default import get_cfg_defaults
+from ASpanFormer.src.utils.misc import lower_config
+
+aspanformer_path = Path(__file__).parent / "../../third_party/ASpanFormer"
+
+
+class ASpanFormer(BaseModel):
+    default_conf = {
+        "model_name": "outdoor.ckpt",
+        "match_threshold": 0.2,
+        "sinkhorn_iterations": 20,
+        "max_keypoints": 2048,
+        "config_path": aspanformer_path / "configs/aspan/outdoor/aspan_test.py",
+    }
+    required_inputs = ["image0", "image1"]
+
+    def _init(self, conf):
+        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"])
+        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)
+        logger.info("Loaded Aspanformer model")
+
+    def _forward(self, data):
+        data_ = {
+            "image0": data["image0"],
+            "image1": data["image1"],
+        }
+        self.net(data_, online_resize=True)
+        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/cotr.py

@@ -0,0 +1,77 @@
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from torchvision.transforms import ToPILImage
+
+from .. import DEVICE, MODEL_REPO_ID
+
+from ..utils.base_model import BaseModel
+
+sys.path.append(str(Path(__file__).parent / "../../third_party/COTR"))
+from COTR.inference.sparse_engine import SparseEngine
+from COTR.models import build_model
+from COTR.options.options import *  # noqa: F403
+from COTR.options.options_utils import *  # noqa: F403
+from COTR.utils import utils as utils_cotr
+
+utils_cotr.fix_randomness(0)
+torch.set_grad_enabled(False)
+
+
+class COTR(BaseModel):
+    default_conf = {
+        "weights": "out/default",
+        "match_threshold": 0.2,
+        "max_keypoints": -1,
+        "model_name": "checkpoint.pth.tar",
+    }
+    required_inputs = ["image0", "image1"]
+
+    def _init(self, conf):
+        parser = argparse.ArgumentParser()
+        set_COTR_arguments(parser)  # noqa: F405
+        opt = parser.parse_args()
+        opt.command = " ".join(sys.argv)
+        opt.load_weights_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+
+        layer_2_channels = {
+            "layer1": 256,
+            "layer2": 512,
+            "layer3": 1024,
+            "layer4": 2048,
+        }
+        opt.dim_feedforward = layer_2_channels[opt.layer]
+
+        model = build_model(opt)
+        model = model.to(DEVICE)
+        weights = torch.load(opt.load_weights_path, map_location="cpu")[
+            "model_state_dict"
+        ]
+        utils_cotr.safe_load_weights(model, weights)
+        self.net = model.eval()
+        self.to_pil_func = ToPILImage(mode="RGB")
+
+    def _forward(self, data):
+        img_a = np.array(self.to_pil_func(data["image0"][0].cpu()))
+        img_b = np.array(self.to_pil_func(data["image1"][0].cpu()))
+        corrs = SparseEngine(
+            self.net, 32, mode="tile"
+        ).cotr_corr_multiscale_with_cycle_consistency(
+            img_a,
+            img_b,
+            np.linspace(0.5, 0.0625, 4),
+            1,
+            max_corrs=self.conf["max_keypoints"],
+            queries_a=None,
+        )
+        pred = {
+            "keypoints0": torch.from_numpy(corrs[:, :2]),
+            "keypoints1": torch.from_numpy(corrs[:, 2:]),
+        }
+        return pred