File size: 2,480 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 |
from kornia.feature.lightglue import LightGlue
from torch import nn
import torch
import os
class LighterGlue(nn.Module):
"""
Lighter version of LightGlue :)
"""
default_conf_xfeat = {
"name": "xfeat", # just for interfacing
"input_dim": 64, # input descriptor dimension (autoselected from weights)
"descriptor_dim": 96,
"add_scale_ori": False,
"add_laf": False, # for KeyNetAffNetHardNet
"scale_coef": 1.0, # to compensate for the SIFT scale bigger than KeyNet
"n_layers": 6,
"num_heads": 1,
"flash": True, # enable FlashAttention if available.
"mp": False, # enable mixed precision
"depth_confidence": -1, # early stopping, disable with -1
"width_confidence": 0.95, # point pruning, disable with -1
"filter_threshold": 0.1, # match threshold
"weights": None,
}
def __init__(self, weights = os.path.abspath(os.path.dirname(__file__)) + '/../weights/xfeat-lighterglue.pt'):
super().__init__()
LightGlue.default_conf = self.default_conf_xfeat
self.net = LightGlue(None)
self.dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if os.path.exists(weights):
state_dict = torch.load(weights, map_location=self.dev)
else:
state_dict = torch.hub.load_state_dict_from_url("https://github.com/verlab/accelerated_features/raw/main/weights/xfeat-lighterglue.pt")
# rename old state dict entries
for i in range(self.net.conf.n_layers):
pattern = f"self_attn.{i}", f"transformers.{i}.self_attn"
state_dict = {k.replace(*pattern): v for k, v in state_dict.items()}
pattern = f"cross_attn.{i}", f"transformers.{i}.cross_attn"
state_dict = {k.replace(*pattern): v for k, v in state_dict.items()}
state_dict = {k.replace('matcher.', ''): v for k, v in state_dict.items()}
self.net.load_state_dict(state_dict, strict=False)
self.net.to(self.dev)
@torch.inference_mode()
def forward(self, data, min_conf = 0.1):
self.net.conf.filter_threshold = min_conf
result = self.net( { 'image0': {'keypoints': data['keypoints0'], 'descriptors': data['descriptors0'], 'image_size': data['image_size0']},
'image1': {'keypoints': data['keypoints1'], 'descriptors': data['descriptors1'], 'image_size': data['image_size1']}
} )
return result
|