import torch
import numpy as np
from omegaconf import OmegaConf
import torchvision.transforms as tfm
from vismatch import get_matcher, BaseMatcher, THIRD_PARTY_DIR
from vismatch.utils import to_numpy, to_tensor, load_module, add_to_path
BASE_PATH = THIRD_PARTY_DIR.joinpath("keypt2subpx")
add_to_path(BASE_PATH)
load_module("gluefactory", BASE_PATH.joinpath("submodules/glue_factory/gluefactory/__init__.py"))
from dataprocess.superpoint_densescore import SuperPoint
add_to_path(THIRD_PARTY_DIR.joinpath("LightGlue"))
from lightglue import LightGlue
from lightglue.utils import rbd, batch_to_device
[docs]
class Keypt2SubpxMatcher(BaseMatcher):
# Maps our detector names to third-party library's expected names
THIRDPARTY_NAMES = {
"superpoint-lightglue": "splg",
"aliked-lightglue": "aliked",
"xfeat": "xfeat",
"xfeat-lightglue": "xfeat",
"dedode": "dedode",
}
def __init__(self, device="cpu", detector_name: str | None = None, **kwargs):
super().__init__(device, **kwargs)
self.detector_name = detector_name
if detector_name == "superpoint-lightglue":
self.matcher = SuperPointDense(self.device)
else:
self.matcher = get_matcher(detector_name, device=device, **kwargs)
detector = self.THIRDPARTY_NAMES[detector_name]
self.keypt2subpx = (
torch.hub.load("KimSinjeong/keypt2subpx", "Keypt2Subpx", pretrained=True, detector=detector, verbose=False)
.eval()
.to(self.device)
)
[docs]
def get_match_idxs(self, mkpts: np.ndarray | torch.Tensor, kpts: np.ndarray | torch.Tensor) -> np.ndarray:
idxs = []
kpts = to_numpy(kpts)
for mkpt in to_numpy(mkpts):
idx = np.flatnonzero(np.all(kpts == mkpt, axis=1)).squeeze().item()
idxs.append(idx)
return np.asarray(idxs)
[docs]
def get_scoremap(self, img=None, idx=None):
assert img is not None or idx is not None, "Must provide either image or idx"
if self.detector_name in ["xfeat", "xfeat-lightglue", "dedode"]:
return None
elif self.detector_name == "aliked-lightglue":
# https://github.com/cvg/LightGlue/blob/edb2b838efb2ecfe3f88097c5fad9887d95aedad/lightglue/aliked.py#L707
return self.matcher.extractor.extract_dense_map(img[None, ...])[-1].squeeze(0)
elif self.detector_name == "superpoint-lightglue":
return self.matcher.get_scoremap(idx)
def _forward(self, img0, img1):
mkpts0, mkpts1, keypoints0, keypoints1, descriptors0, descriptors1 = self.matcher._forward(img0, img1)
if len(mkpts0): # only run subpx refinement if kpts are found
matching_idxs0, matching_idxs1 = (
self.get_match_idxs(mkpts0, keypoints0),
self.get_match_idxs(mkpts1, keypoints1),
)
mdesc0, mdesc1 = descriptors0[matching_idxs0], descriptors1[matching_idxs1]
scores0, scores1 = self.get_scoremap(img0, 0), self.get_scoremap(img1, 1)
mkpts0, mkpts1 = self.keypt2subpx(
to_tensor(mkpts0, self.device),
to_tensor(mkpts1, self.device),
img0,
img1,
mdesc0,
mdesc1,
scores0,
scores1,
)
return mkpts0, mkpts1, keypoints0, keypoints1, descriptors0, descriptors1
[docs]
class SuperPointDense(BaseMatcher):
# SuperPoint, with Dense Scoremap for Keypt2Subpx refinement
modelconf = {
"name": "two_view_pipeline",
"extractor": {
"name": "superpoint_densescore",
"max_num_keypoints": 2048,
"force_num_keypoints": False,
"detection_threshold": 0.0,
"nms_radius": 3,
"remove_borders": 3,
"trainable": False,
},
"matcher": {
"name": "matchers.lightglue_wrapper",
"weights": "superpoint",
"depth_confidence": -1,
"width_confidence": -1,
"filter_threshold": 0.1,
"trainable": False,
},
"ground_truth": {"name": "matchers.depth_matcher", "th_positive": 3, "th_negative": 5, "th_epi": 5},
"allow_no_extract": True,
}
def __init__(self, device="cpu", **kwargs):
super().__init__(device, **kwargs)
self.config = OmegaConf.create(self.modelconf)
self.extractor = SuperPoint(self.config).to(self.device).eval()
self.matcher = LightGlue(features="superpoint", depth_confidence=-1, width_confidence=-1).to(self.device)
self.scoremaps = {}
[docs]
def preprocess(self, img):
return tfm.Grayscale()(img).unsqueeze(0)
[docs]
def get_scoremap(self, idx):
return self.scoremaps[idx]
def _forward(self, img0, img1):
img0 = self.preprocess(img0)
img1 = self.preprocess(img1)
feats0 = self.extractor({"image": img0})
feats1 = self.extractor({"image": img1})
self.scoremaps[0] = feats0["keypoint_scores"]
self.scoremaps[1] = feats1["keypoint_scores"]
# requires keys ['keypoints', 'keypoint_scores', 'descriptors', 'image_size']
matches01 = self.matcher({"image0": feats0, "image1": feats1})
data = [feats0, feats1, matches01]
# remove batch dim and move to target device
feats0, feats1, matches01 = [batch_to_device(rbd(x), self.device) for x in data]
kpts0, kpts1, matches = (
feats0["keypoints"],
feats1["keypoints"],
matches01["matches"],
)
desc0 = feats0["descriptors"]
desc1 = feats1["descriptors"]
mkpts0, mkpts1 = kpts0[matches[..., 0]], kpts1[matches[..., 1]]
return mkpts0, mkpts1, kpts0, kpts1, desc0, desc1