rtabmap_superglue.py
Go to the documentation of this file.
1 #! /usr/bin/env python3
2 #
3 # Drop this file in the root folder of SuperGlue git: https://github.com/magicleap/SuperGluePretrainedNetwork
4 # To use with rtabmap:
5 # --Vis/CorNNType 6 --SuperGlue/Path "~/SuperGluePretrainedNetwork/rtabmap_superglue.py"
6 #
7 
8 import random
9 import numpy as np
10 import torch
11 
12 #import sys
13 #import os
14 #print(os.sys.path)
15 #print(sys.version)
16 
17 from models.matching import SuperGlue
18 
19 torch.set_grad_enabled(False)
20 
21 device = 'cpu'
22 superglue = []
23 
24 def init(descriptorDim, matchThreshold, iterations, cuda, model):
25  print("SuperGlue python init()")
26  # Load the SuperGlue model.
27  global device
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'"
30  config = {
31  'superglue': {
32  'weights': model,
33  'sinkhorn_iterations': iterations,
34  'match_threshold': matchThreshold,
35  'descriptor_dim' : descriptorDim
36  }
37  }
38  global superglue
39  superglue = SuperGlue(config.get('superglue', {})).eval().to(device)
40 
41 
42 def match(kptsFrom, kptsTo, scoresFrom, scoresTo, descriptorsFrom, descriptorsTo, imageWidth, imageHeight):
43  #print("SuperGlue python match()")
44  global device
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, :, :]
57 
58  data = {
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),
67  }
68 
69 
70  global superglue
71  results = superglue(data)
72 
73  matches0 = results['matches0'].to('cpu').numpy()
74 
75  matchesFrom = np.nonzero(matches0!=-1)[1]
76  matchesTo = matches0[np.nonzero(matches0!=-1)]
77 
78  matchesArray = np.stack((matchesFrom, matchesTo), axis=1);
79 
80  return matchesArray
81 
82 
83 if __name__ == '__main__':
84  #test
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)
def match(kptsFrom, kptsTo, scoresFrom, scoresTo, descriptorsFrom, descriptorsTo, imageWidth, imageHeight)
def init(descriptorDim, matchThreshold, iterations, cuda, model)


rtabmap
Author(s): Mathieu Labbe
autogenerated on Mon Jan 23 2023 03:37:30