File size: 3,148 Bytes
0a82b18 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from matching import BaseMatcher, THIRD_PARTY_DIR
from matching.utils import add_to_path
from lightglue import match_pair
from lightglue import LightGlue, SuperPoint, DISK, SIFT, ALIKED, DoGHardNet
add_to_path(THIRD_PARTY_DIR.joinpath("LightGlue"))
class LightGlueBase(BaseMatcher):
"""
This class is the parent for all methods that use LightGlue as a matcher,
with different local features. It implements the forward which is the same
regardless of the feature extractor of choice.
Therefore this class should *NOT* be instatiated, as it needs its children to define
the extractor and the matcher.
"""
def __init__(self, device="cpu", **kwargs):
super().__init__(device, **kwargs)
def _forward(self, img0, img1):
"""
"extractor" and "matcher" are instantiated by the subclasses.
"""
feats0, feats1, matches01 = match_pair(self.extractor, self.matcher, img0, img1, device=self.device)
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
class SiftLightGlue(LightGlueBase):
def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
super().__init__(device, **kwargs)
self.extractor = SIFT(max_num_keypoints=max_num_keypoints).eval().to(self.device)
self.matcher = LightGlue(features="sift", depth_confidence=-1, width_confidence=-1).to(self.device)
class SuperpointLightGlue(LightGlueBase):
def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
super().__init__(device, **kwargs)
self.extractor = SuperPoint(max_num_keypoints=max_num_keypoints).eval().to(self.device)
self.matcher = LightGlue(features="superpoint", depth_confidence=-1, width_confidence=-1).to(self.device)
class DiskLightGlue(LightGlueBase):
def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
super().__init__(device, **kwargs)
self.extractor = DISK(max_num_keypoints=max_num_keypoints).eval().to(self.device)
self.matcher = LightGlue(features="disk", depth_confidence=-1, width_confidence=-1).to(self.device)
class AlikedLightGlue(LightGlueBase):
def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
super().__init__(device, **kwargs)
self.extractor = ALIKED(max_num_keypoints=max_num_keypoints).eval().to(self.device)
self.matcher = LightGlue(features="aliked", depth_confidence=-1, width_confidence=-1).to(self.device)
class DognetLightGlue(LightGlueBase):
def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
super().__init__(device, **kwargs)
self.extractor = DoGHardNet(max_num_keypoints=max_num_keypoints).eval().to(self.device)
self.matcher = LightGlue(features="doghardnet", depth_confidence=-1, width_confidence=-1).to(self.device)
|