17 from models.matching
import SuperGlue
19 torch.set_grad_enabled(
False)
24 def init(descriptorDim, matchThreshold, iterations, cuda, model):
25 print(
"SuperGlue python init()")
28 device =
'cuda' if torch.cuda.is_available()
and cuda
else 'cpu'
29 assert model ==
"indoor" or model ==
"outdoor",
"Available models for SuperGlue are 'indoor' or 'outdoor'"
33 'sinkhorn_iterations': iterations,
34 'match_threshold': matchThreshold,
35 'descriptor_dim' : descriptorDim
39 superglue = SuperGlue(config.get(
'superglue', {})).eval().to(device)
42 def match(kptsFrom, kptsTo, scoresFrom, scoresTo, descriptorsFrom, descriptorsTo, imageWidth, imageHeight):
45 kptsFrom = np.asarray(kptsFrom)
46 kptsFrom = kptsFrom[
None, :, :]
47 kptsTo = np.asarray(kptsTo)
48 kptsTo = kptsTo[
None, :, :]
49 scoresFrom = np.asarray(scoresFrom)
50 scoresFrom = scoresFrom[
None, :]
51 scoresTo = np.asarray(scoresTo)
52 scoresTo = scoresTo[
None, :]
53 descriptorsFrom = np.transpose(np.asarray(descriptorsFrom))
54 descriptorsFrom = descriptorsFrom[
None, :, :]
55 descriptorsTo = np.transpose(np.asarray(descriptorsTo))
56 descriptorsTo = descriptorsTo[
None, :, :]
59 'image0': torch.rand(1, 1, imageHeight, imageWidth).to(device),
60 'image1': torch.rand(1, 1, imageHeight, imageWidth).to(device),
61 'scores0': torch.from_numpy(scoresFrom).to(device),
62 'scores1': torch.from_numpy(scoresTo).to(device),
63 'keypoints0': torch.from_numpy(kptsFrom).to(device),
64 'keypoints1': torch.from_numpy(kptsTo).to(device),
65 'descriptors0': torch.from_numpy(descriptorsFrom).to(device),
66 'descriptors1': torch.from_numpy(descriptorsTo).to(device),
73 matches0 = results[
'matches0'].to(
'cpu').numpy()
75 matchesFrom = np.nonzero(matches0!=-1)[1]
76 matchesTo = matches0[np.nonzero(matches0!=-1)]
78 matchesArray = np.stack((matchesFrom, matchesTo), axis=1);
83 if __name__ ==
'__main__':
85 init(256, 0.2, 20,
True,
'indoor')
86 match([[1, 2], [1,3]], [[1, 3], [1,2]], [1, 3], [1,3], np.full((2, 256), 1),np.full((2, 256), 1), 640, 480)