Vincentqyw
fix: roma
8b973ee
raw
history blame
1.33 kB
import torch
from .nn import NN2
from darkfeat import DarkFeat
class NNMatching(torch.nn.Module):
def __init__(self, model_path=""):
super().__init__()
self.nn = NN2().eval()
self.darkfeat = DarkFeat(model_path).eval()
def forward(self, data):
"""Run DarkFeat and nearest neighborhood matching
Args:
data: dictionary with minimal keys: ['image0', 'image1']
"""
pred = {}
# Extract DarkFeat (keypoints, scores, descriptors)
if "keypoints0" not in data:
pred0 = self.darkfeat({"image": data["image0"]})
# print({k+'0': v[0].shape for k, v in pred0.items()})
pred = {**pred, **{k + "0": [v] for k, v in pred0.items()}}
if "keypoints1" not in data:
pred1 = self.darkfeat({"image": data["image1"]})
pred = {**pred, **{k + "1": [v] for k, v in pred1.items()}}
# Batch all features
# We should either have i) one image per batch, or
# ii) the same number of local features for all images in the batch.
data = {**data, **pred}
for k in data:
if isinstance(data[k], (list, tuple)):
data[k] = torch.stack(data[k])
# Perform the matching
pred = {**pred, **self.nn(data)}
return pred