vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/MatchAnything/imcui/hloc/extractors/cosplace.py

@@ -0,0 +1,44 @@
+"""
+Code for loading models trained with CosPlace as a global features extractor
+for geolocalization through image retrieval.
+Multiple models are available with different backbones. Below is a summary of
+models available (backbone : list of available output descriptors
+dimensionality). For example you can use a model based on a ResNet50 with
+descriptors dimensionality 1024.
+    ResNet18:  [32, 64, 128, 256, 512]
+    ResNet50:  [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet101: [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet152: [32, 64, 128, 256, 512, 1024, 2048]
+    VGG16:     [    64, 128, 256, 512]
+
+CosPlace paper: https://arxiv.org/abs/2204.02287
+"""
+
+import torch
+import torchvision.transforms as tvf
+
+from ..utils.base_model import BaseModel
+
+
+class CosPlace(BaseModel):
+    default_conf = {"backbone": "ResNet50", "fc_output_dim": 2048}
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "gmberton/CosPlace",
+            "get_trained_model",
+            backbone=conf["backbone"],
+            fc_output_dim=conf["fc_output_dim"],
+        ).eval()
+
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        self.norm_rgb = tvf.Normalize(mean=mean, std=std)
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+        desc = self.net(image)
+        return {
+            "global_descriptor": desc,
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/d2net.py

@@ -0,0 +1,60 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import MODEL_REPO_ID, logger
+from ..utils.base_model import BaseModel
+
+d2net_path = Path(__file__).parent / "../../third_party/d2net"
+sys.path.append(str(d2net_path))
+from lib.model_test import D2Net as _D2Net
+from lib.pyramid import process_multiscale
+
+
+class D2Net(BaseModel):
+    default_conf = {
+        "model_name": "d2_tf.pth",
+        "checkpoint_dir": d2net_path / "models",
+        "use_relu": True,
+        "multiscale": False,
+        "max_keypoints": 1024,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        logger.info("Loading D2Net model...")
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
+        )
+        logger.info(f"Loading model from {model_path}...")
+        self.net = _D2Net(
+            model_file=model_path, use_relu=conf["use_relu"], use_cuda=False
+        )
+        logger.info("Load D2Net model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        image = image.flip(1)  # RGB -> BGR
+        norm = image.new_tensor([103.939, 116.779, 123.68])
+        image = image * 255 - norm.view(1, 3, 1, 1)  # caffe normalization
+
+        if self.conf["multiscale"]:
+            keypoints, scores, descriptors = process_multiscale(image, self.net)
+        else:
+            keypoints, scores, descriptors = process_multiscale(
+                image, self.net, scales=[1]
+            )
+        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],
+            "descriptors": torch.from_numpy(descriptors.T)[None],
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/darkfeat.py

@@ -0,0 +1,44 @@
+import sys
+from pathlib import Path
+
+from .. import MODEL_REPO_ID, logger
+
+from ..utils.base_model import BaseModel
+
+darkfeat_path = Path(__file__).parent / "../../third_party/DarkFeat"
+sys.path.append(str(darkfeat_path))
+from darkfeat import DarkFeat as DarkFeat_
+
+
+class DarkFeat(BaseModel):
+    default_conf = {
+        "model_name": "DarkFeat.pth",
+        "max_keypoints": 1000,
+        "detection_threshold": 0.5,
+        "sub_pixel": False,
+    }
+    required_inputs = ["image"]
+
+    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("Loaded DarkFeat model: {}".format(model_path))
+        self.net = DarkFeat_(model_path)
+        logger.info("Load DarkFeat model done.")
+
+    def _forward(self, data):
+        pred = self.net({"image": data["image"]})
+        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
+            "descriptors": descriptors[None],  # 1 x 128 x N
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/dedode.py

@@ -0,0 +1,86 @@
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.transforms as transforms
+
+from .. import MODEL_REPO_ID, logger
+
+from ..utils.base_model import BaseModel
+
+dedode_path = Path(__file__).parent / "../../third_party/DeDoDe"
+sys.path.append(str(dedode_path))
+
+from DeDoDe import dedode_descriptor_B, dedode_detector_L
+from DeDoDe.utils import to_pixel_coords
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class DeDoDe(BaseModel):
+    default_conf = {
+        "name": "dedode",
+        "model_detector_name": "dedode_detector_L.pth",
+        "model_descriptor_name": "dedode_descriptor_B.pth",
+        "max_keypoints": 2000,
+        "match_threshold": 0.2,
+        "dense": False,  # Now fixed to be false
+    }
+    required_inputs = [
+        "image",
+    ]
+
+    # Initialize the line matcher
+    def _init(self, conf):
+        model_detector_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, conf["model_detector_name"]),
+        )
+        model_descriptor_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(Path(__file__).stem, conf["model_descriptor_name"]),
+        )
+        logger.info("Loaded DarkFeat model: {}".format(model_detector_path))
+        self.normalizer = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+        # 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)
+        logger.info("Load DeDoDe model done.")
+
+    def _forward(self, data):
+        """
+        data: dict, keys: {'image0','image1'}
+        image shape: N x C x H x W
+        color mode: RGB
+        """
+        img0 = self.normalizer(data["image"].squeeze()).float()[None]
+        H_A, W_A = img0.shape[2:]
+
+        # step 1: detect keypoints
+        detections_A = None
+        batch_A = {"image": img0}
+        if self.conf["dense"]:
+            detections_A = self.detector.detect_dense(batch_A)
+        else:
+            detections_A = self.detector.detect(
+                batch_A, num_keypoints=self.conf["max_keypoints"]
+            )
+        keypoints_A, P_A = detections_A["keypoints"], detections_A["confidence"]
+
+        # step 2: describe keypoints
+        # dim: 1 x N x 256
+        description_A = self.descriptor.describe_keypoints(batch_A, keypoints_A)[
+            "descriptions"
+        ]
+        keypoints_A = to_pixel_coords(keypoints_A, H_A, W_A)
+
+        return {
+            "keypoints": keypoints_A,  # 1 x N x 2
+            "descriptors": description_A.permute(0, 2, 1),  # 1 x 256 x N
+            "scores": P_A,  # 1 x N
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/dir.py

@@ -0,0 +1,78 @@
+import os
+import sys
+from pathlib import Path
+from zipfile import ZipFile
+
+import gdown
+import sklearn
+import torch
+
+from ..utils.base_model import BaseModel
+
+sys.path.append(str(Path(__file__).parent / "../../third_party/deep-image-retrieval"))
+os.environ["DB_ROOT"] = ""  # required by dirtorch
+
+from dirtorch.extract_features import load_model  # noqa: E402
+from dirtorch.utils import common  # noqa: E402
+
+# The DIR model checkpoints (pickle files) include sklearn.decomposition.pca,
+# which has been deprecated in sklearn v0.24
+# and must be explicitly imported with `from sklearn.decomposition import PCA`.
+# This is a hacky workaround to maintain forward compatibility.
+sys.modules["sklearn.decomposition.pca"] = sklearn.decomposition._pca
+
+
+class DIR(BaseModel):
+    default_conf = {
+        "model_name": "Resnet-101-AP-GeM",
+        "whiten_name": "Landmarks_clean",
+        "whiten_params": {
+            "whitenp": 0.25,
+            "whitenv": None,
+            "whitenm": 1.0,
+        },
+        "pooling": "gem",
+        "gemp": 3,
+    }
+    required_inputs = ["image"]
+
+    dir_models = {
+        "Resnet-101-AP-GeM": "https://docs.google.com/uc?export=download&id=1UWJGDuHtzaQdFhSMojoYVQjmCXhIwVvy",
+    }
+
+    def _init(self, conf):
+        # todo: download from google drive -> huggingface models
+        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"]]
+            gdown.download(str(link), str(checkpoint) + ".zip", quiet=False)
+            zf = ZipFile(str(checkpoint) + ".zip", "r")
+            zf.extractall(checkpoint.parent)
+            zf.close()
+            os.remove(str(checkpoint) + ".zip")
+
+        self.net = load_model(checkpoint, False)  # first load on CPU
+        if conf["whiten_name"]:
+            assert conf["whiten_name"] in self.net.pca
+
+    def _forward(self, data):
+        image = data["image"]
+        assert image.shape[1] == 3
+        mean = self.net.preprocess["mean"]
+        std = self.net.preprocess["std"]
+        image = image - image.new_tensor(mean)[:, None, None]
+        image = image / image.new_tensor(std)[:, None, None]
+
+        desc = self.net(image)
+        desc = desc.unsqueeze(0)  # batch dimension
+        if self.conf["whiten_name"]:
+            pca = self.net.pca[self.conf["whiten_name"]]
+            desc = common.whiten_features(
+                desc.cpu().numpy(), pca, **self.conf["whiten_params"]
+            )
+            desc = torch.from_numpy(desc)
+
+        return {
+            "global_descriptor": desc,
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/disk.py

@@ -0,0 +1,35 @@
+import kornia
+
+from .. import logger
+
+from ..utils.base_model import BaseModel
+
+
+class DISK(BaseModel):
+    default_conf = {
+        "weights": "depth",
+        "max_keypoints": None,
+        "nms_window_size": 5,
+        "detection_threshold": 0.0,
+        "pad_if_not_divisible": True,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.model = kornia.feature.DISK.from_pretrained(conf["weights"])
+        logger.info("Load DISK model done.")
+
+    def _forward(self, data):
+        image = data["image"]
+        features = self.model(
+            image,
+            n=self.conf["max_keypoints"],
+            window_size=self.conf["nms_window_size"],
+            score_threshold=self.conf["detection_threshold"],
+            pad_if_not_divisible=self.conf["pad_if_not_divisible"],
+        )
+        return {
+            "keypoints": [f.keypoints for f in features][0][None],
+            "scores": [f.detection_scores for f in features][0][None],
+            "descriptors": [f.descriptors.t() for f in features][0][None],
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/dog.py

@@ -0,0 +1,135 @@
+import kornia
+import numpy as np
+import pycolmap
+import torch
+from kornia.feature.laf import (
+    extract_patches_from_pyramid,
+    laf_from_center_scale_ori,
+)
+
+from ..utils.base_model import BaseModel
+
+EPS = 1e-6
+
+
+def sift_to_rootsift(x):
+    x = x / (np.linalg.norm(x, ord=1, axis=-1, keepdims=True) + EPS)
+    x = np.sqrt(x.clip(min=EPS))
+    x = x / (np.linalg.norm(x, axis=-1, keepdims=True) + EPS)
+    return x
+
+
+class DoG(BaseModel):
+    default_conf = {
+        "options": {
+            "first_octave": 0,
+            "peak_threshold": 0.01,
+        },
+        "descriptor": "rootsift",
+        "max_keypoints": -1,
+        "patch_size": 32,
+        "mr_size": 12,
+    }
+    required_inputs = ["image"]
+    detection_noise = 1.0
+    max_batch_size = 1024
+
+    def _init(self, conf):
+        if conf["descriptor"] == "sosnet":
+            self.describe = kornia.feature.SOSNet(pretrained=True)
+        elif conf["descriptor"] == "hardnet":
+            self.describe = kornia.feature.HardNet(pretrained=True)
+        elif conf["descriptor"] not in ["sift", "rootsift"]:
+            raise ValueError(f'Unknown descriptor: {conf["descriptor"]}')
+
+        self.sift = None  # lazily instantiated on the first image
+        self.dummy_param = torch.nn.Parameter(torch.empty(0))
+        self.device = torch.device("cpu")
+
+    def to(self, *args, **kwargs):
+        device = kwargs.get("device")
+        if device is None:
+            match = [a for a in args if isinstance(a, (torch.device, str))]
+            if len(match) > 0:
+                device = match[0]
+        if device is not None:
+            self.device = torch.device(device)
+        return super().to(*args, **kwargs)
+
+    def _forward(self, data):
+        image = data["image"]
+        image_np = image.cpu().numpy()[0, 0]
+        assert image.shape[1] == 1
+        assert image_np.min() >= -EPS and image_np.max() <= 1 + EPS
+
+        if self.sift is None:
+            device = self.dummy_param.device
+            use_gpu = pycolmap.has_cuda and device.type == "cuda"
+            options = {**self.conf["options"]}
+            if self.conf["descriptor"] == "rootsift":
+                options["normalization"] = pycolmap.Normalization.L1_ROOT
+            else:
+                options["normalization"] = pycolmap.Normalization.L2
+            self.sift = pycolmap.Sift(
+                options=pycolmap.SiftExtractionOptions(options),
+                device=getattr(pycolmap.Device, "cuda" if use_gpu else "cpu"),
+            )
+        keypoints, descriptors = self.sift.extract(image_np)
+        scales = keypoints[:, 2]
+        oris = np.rad2deg(keypoints[:, 3])
+
+        if self.conf["descriptor"] in ["sift", "rootsift"]:
+            # We still renormalize because COLMAP does not normalize well,
+            # maybe due to numerical errors
+            if self.conf["descriptor"] == "rootsift":
+                descriptors = sift_to_rootsift(descriptors)
+            descriptors = torch.from_numpy(descriptors)
+        elif self.conf["descriptor"] in ("sosnet", "hardnet"):
+            center = keypoints[:, :2] + 0.5
+            laf_scale = scales * self.conf["mr_size"] / 2
+            laf_ori = -oris
+            lafs = laf_from_center_scale_ori(
+                torch.from_numpy(center)[None],
+                torch.from_numpy(laf_scale)[None, :, None, None],
+                torch.from_numpy(laf_ori)[None, :, None],
+            ).to(image.device)
+            patches = extract_patches_from_pyramid(
+                image, lafs, PS=self.conf["patch_size"]
+            )[0]
+            descriptors = patches.new_zeros((len(patches), 128))
+            if len(patches) > 0:
+                for start_idx in range(0, len(patches), self.max_batch_size):
+                    end_idx = min(len(patches), start_idx + self.max_batch_size)
+                    descriptors[start_idx:end_idx] = self.describe(
+                        patches[start_idx:end_idx]
+                    )
+        else:
+            raise ValueError(f'Unknown descriptor: {self.conf["descriptor"]}')
+
+        keypoints = torch.from_numpy(keypoints[:, :2])  # keep only x, y
+        scales = torch.from_numpy(scales)
+        oris = torch.from_numpy(oris)
+        scores = keypoints.new_zeros(len(keypoints))  # no scores for SIFT yet
+
+        if self.conf["max_keypoints"] != -1:
+            # TODO: check that the scores from PyCOLMAP are 100% correct,
+            # follow https://github.com/mihaidusmanu/pycolmap/issues/8
+            max_number = (
+                scores.shape[0]
+                if scores.shape[0] < self.conf["max_keypoints"]
+                else self.conf["max_keypoints"]
+            )
+            values, indices = torch.topk(scores, max_number)
+            keypoints = keypoints[indices]
+            scales = scales[indices]
+            oris = oris[indices]
+            scores = scores[indices]
+            descriptors = descriptors[indices]
+
+        return {
+            "keypoints": keypoints[None],
+            "scales": scales[None],
+            "oris": oris[None],
+            "scores": scores[None],
+            "descriptors": descriptors.T[None],
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/eigenplaces.py

@@ -0,0 +1,57 @@
+"""
+Code for loading models trained with EigenPlaces (or CosPlace) as a global
+features extractor for geolocalization through image retrieval.
+Multiple models are available with different backbones. Below is a summary of
+models available (backbone : list of available output descriptors
+dimensionality). For example you can use a model based on a ResNet50 with
+descriptors dimensionality 1024.
+
+EigenPlaces trained models:
+    ResNet18:  [     256, 512]
+    ResNet50:  [128, 256, 512, 2048]
+    ResNet101: [128, 256, 512, 2048]
+    VGG16:     [     512]
+
+CosPlace trained models:
+    ResNet18:  [32, 64, 128, 256, 512]
+    ResNet50:  [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet101: [32, 64, 128, 256, 512, 1024, 2048]
+    ResNet152: [32, 64, 128, 256, 512, 1024, 2048]
+    VGG16:     [    64, 128, 256, 512]
+
+EigenPlaces paper (ICCV 2023): https://arxiv.org/abs/2308.10832
+CosPlace paper (CVPR 2022): https://arxiv.org/abs/2204.02287
+"""
+
+import torch
+import torchvision.transforms as tvf
+
+from ..utils.base_model import BaseModel
+
+
+class EigenPlaces(BaseModel):
+    default_conf = {
+        "variant": "EigenPlaces",
+        "backbone": "ResNet101",
+        "fc_output_dim": 2048,
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        self.net = torch.hub.load(
+            "gmberton/" + conf["variant"],
+            "get_trained_model",
+            backbone=conf["backbone"],
+            fc_output_dim=conf["fc_output_dim"],
+        ).eval()
+
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        self.norm_rgb = tvf.Normalize(mean=mean, std=std)
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+        desc = self.net(image)
+        return {
+            "global_descriptor": desc,
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/example.py

@@ -0,0 +1,56 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from .. import logger
+from ..utils.base_model import BaseModel
+
+example_path = Path(__file__).parent / "../../third_party/example"
+sys.path.append(str(example_path))
+
+# import some modules here
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Example(BaseModel):
+    # change to your default configs
+    default_conf = {
+        "name": "example",
+        "keypoint_threshold": 0.1,
+        "max_keypoints": 2000,
+        "model_name": "model.pth",
+    }
+    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():
+            logger.info(f"No model found at {model_path}")
+
+        # init model
+        self.net = callable
+        # self.net = ExampleNet(is_test=True)
+        state_dict = torch.load(model_path, map_location="cpu")
+        self.net.load_state_dict(state_dict["model_state"])
+        logger.info("Load example model done.")
+
+    def _forward(self, data):
+        # data: dict, keys: 'image'
+        # image color mode: RGB
+        # image value range in [0, 1]
+        image = data["image"]
+
+        # B: batch size, N: number of keypoints
+        # keypoints shape: B x N x 2, type: torch tensor
+        # scores shape: B x N, type: torch tensor
+        # descriptors shape: B x 128 x N, type: torch tensor
+        keypoints, scores, descriptors = self.net(image)
+
+        return {
+            "keypoints": keypoints,
+            "scores": scores,
+            "descriptors": descriptors,
+        }

vismatch/third_party/MatchAnything/imcui/hloc/extractors/fire.py

@@ -0,0 +1,72 @@
+import logging
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.transforms as tvf
+
+from ..utils.base_model import BaseModel
+
+logger = logging.getLogger(__name__)
+fire_path = Path(__file__).parent / "../../third_party/fire"
+sys.path.append(str(fire_path))
+
+
+import fire_network
+
+
+class FIRe(BaseModel):
+    default_conf = {
+        "global": True,
+        "asmk": False,
+        "model_name": "fire_SfM_120k.pth",
+        "scales": [2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25],  # default params
+        "features_num": 1000,  # TODO:not supported now
+        "asmk_name": "asmk_codebook.bin",  # TODO:not supported now
+        "config_name": "eval_fire.yml",
+    }
+    required_inputs = ["image"]
+
+    # Models exported using
+    fire_models = {
+        "fire_SfM_120k.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/official/fire.pth",
+        "fire_imagenet.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/pretraining/fire_imagenet.pth",
+    }
+
+    def _init(self, conf):
+        assert conf["model_name"] in self.fire_models.keys()
+        # Config paths
+        model_path = fire_path / "model" / conf["model_name"]
+
+        # Download the model.
+        if not model_path.exists():
+            model_path.parent.mkdir(exist_ok=True)
+            link = self.fire_models[conf["model_name"]]
+            cmd = ["wget", "--quiet", link, "-O", str(model_path)]
+            logger.info(f"Downloading the FIRe model with `{cmd}`.")
+            subprocess.run(cmd, check=True)
+
+        logger.info("Loading fire model...")
+
+        # Load net
+        state = torch.load(model_path)
+        state["net_params"]["pretrained"] = None
+        net = fire_network.init_network(**state["net_params"])
+        net.load_state_dict(state["state_dict"])
+        self.net = net
+
+        self.norm_rgb = tvf.Normalize(
+            **dict(zip(["mean", "std"], net.runtime["mean_std"]))
+        )
+
+        # params
+        self.scales = conf["scales"]
+
+    def _forward(self, data):
+        image = self.norm_rgb(data["image"])
+
+        # Feature extraction.
+        desc = self.net.forward_global(image, scales=self.scales)
+
+        return {"global_descriptor": desc}