Fixed SG endpoint matching when there are no lines (cvg#7)

Author: rpautrat

limap/line2d/endpoints/extractor.py

@@ -9,6 +9,7 @@
 sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 import limap.util.io as limapio
 
+
 class SuperPointEndpointsExtractor(BaseDetector):
     def __init__(self, options = BaseDetectorOptions(), device=None):
         super(SuperPointEndpointsExtractor, self).__init__(options)
@@ -45,6 +46,11 @@ def compute_descinfo(self, img, segs):
             - the descriptor of each endpoints of shape [256, N*2]
         """
         lines = segs[:, :4].reshape(-1, 2)
+        if len(lines) == 0:
+            return {
+                'image_shape': img.shape, 'lines': lines,
+                'lines_score': np.zeros((0,)),
+                'endpoints_desc': np.zeros((256, 0))}
         scores = segs[:, -1] * np.sqrt(np.linalg.norm(segs[:, :2]
                                                       - segs[:, 2:4], axis=1))
         scores /= np.amax(scores) + 1e-8
@@ -58,5 +64,3 @@ def compute_descinfo(self, img, segs):
                 torch_img, torch_endpoints)['descriptors'][0].cpu().numpy()
         return {'image_shape': img.shape, 'lines': lines,
                 'lines_score': scores, 'endpoints_desc': endpoint_descs}
-
-

limap/line2d/endpoints/matcher.py

@@ -9,6 +9,7 @@
 sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 import limap.util.io as limapio
 
+
 class NNEndpointsMatcher(BaseMatcher):
     def __init__(self, extractor, options = BaseMatcherOptions(), device=None):
         super(NNEndpointsMatcher, self).__init__(extractor, options)
@@ -87,6 +88,7 @@ def match_segs_with_descinfo_topk(self, descinfo1, descinfo2, topk=10):
                               matches.flatten()], axis=1)
         return matches_t
 
+
 class SuperGlueEndpointsMatcher(BaseMatcher):
     def __init__(self, extractor, options = BaseMatcherOptions(),
                  weights='outdoor', device=None):
@@ -99,10 +101,13 @@ def get_module_name(self):
         return "superglue_endpoints"
 
     def match_pair(self, descinfo1, descinfo2):
+        if len(descinfo1['lines']) == 0 or len(descinfo1['lines']) == 0:
+            return np.empty((0, 2))
         if self.topk == 0:
             return self.match_segs_with_descinfo(descinfo1, descinfo2)
         else:
-            return self.match_segs_with_descinfo_topk(descinfo1, descinfo2, topk=self.topk)
+            return self.match_segs_with_descinfo_topk(descinfo1, descinfo2,
+                                                      topk=self.topk)
 
     def match_segs_with_descinfo(self, descinfo1, descinfo2):
         # Setup the inputs for SuperGlue
@@ -199,4 +204,3 @@ def match_segs_with_descinfo_topk(self, descinfo1, descinfo2, topk=10):
         matches_t = np.stack([np.arange(n_lines).repeat(topk),
                               matches.flatten()], axis=1)
         return matches_t
-