Add scalar detection uncertainty (#158)

Author: skydes

hloc/extract_features.py

@@ -241,6 +241,8 @@ def main(conf: Dict,
             size = np.array(data['image'].shape[-2:][::-1])
             scales = (original_size / size).astype(np.float32)
             pred['keypoints'] = (pred['keypoints'] + .5) * scales[None] - .5
+            # add keypoint uncertainties scaled to the original resolution
+            uncertainty = getattr(model, 'detection_noise', 1) * scales.mean()
 
         if as_half:
             for k in pred:
@@ -255,6 +257,8 @@ def main(conf: Dict,
                 grp = fd.create_group(name)
                 for k, v in pred.items():
                     grp.create_dataset(k, data=v)
+                if 'keypoints' in pred:
+                    grp['keypoints'].attrs['uncertainty'] = uncertainty
             except OSError as error:
                 if 'No space left on device' in error.args[0]:
                     logger.error(

hloc/extractors/sift.py

@@ -29,6 +29,7 @@ class SIFT(BaseModel):
         'max_keypoints': -1
     }
     required_inputs = ['image']
+    detection_noise = 1.0
 
     def _init(self, conf):
         self.root = conf['root']

hloc/extractors/superpoint.py

@@ -53,6 +53,7 @@ class SuperPoint(BaseModel):
         'fix_sampling': False,
     }
     required_inputs = ['image']
+    detection_noise = 2.0
 
     def _init(self, conf):
         if conf['fix_sampling']: