Spaces:
Running
Running
| 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 | |