Fix dense assignment without aggregation (#243)

hloc/match_dense.py

@@ -87,8 +87,8 @@
 
 def to_cpts(kpts, ps):
     if ps > 0.0:
-        cpts = np.round(np.round((kpts + 0.5) / ps) * ps - 0.5, 2)
-    return [tuple(cpt) for cpt in cpts]
+        kpts = np.round(np.round((kpts + 0.5) / ps) * ps - 0.5, 2)
+    return [tuple(cpt) for cpt in kpts]
 
 
 def assign_keypoints(kpts: np.ndarray,
@@ -106,7 +106,7 @@ def assign_keypoints(kpts: np.ndarray,
         return kpt_ids
     else:
         ps = cell_size if cell_size is not None else max_error
-        ps = max(cell_size, max_error)
+        ps = max(ps, max_error)
         # With update we quantize and bin (optionally)
         assert isinstance(other_cpts, list)
         kpt_ids = []
@@ -270,13 +270,15 @@ def match_dense(conf: Dict,
             image0, image1 = image0.to(device), image1.to(device)
 
             # match semi-dense
-            if name1 in existing_refs:
-                # flip to enable refinement in query image
+            # for consistency with pairs_from_*: refine kpts of image0
+            if name0 in existing_refs:
+                # special case: flip to enable refinement in query image
                 pred = model({'image0': image1, 'image1': image0})
                 pred = {**pred,
                         'keypoints0': pred['keypoints1'],
                         'keypoints1': pred['keypoints0']}
             else:
+                # usual case
                 pred = model({'image0': image0, 'image1': image1})
 
             # Rescale keypoints and move to cpu
@@ -376,16 +378,18 @@ def aggregate_matches(
             update1 = name1 in required_queries
 
             # in localization we do not want to bin the query kp
-            if update1 and not update0 and max_kps is None:
-                max_error1 = 0.0
+            # assumes that the query is name0!
+            if update0 and not update1 and max_kps is None:
+                max_error0 = cell_size0 = 0.0
             else:
-                max_error1 = conf['max_error']
+                max_error0 = conf['max_error']
+                cell_size0 = conf['cell_size']
 
             # Get match ids and extend query keypoints (cpdict)
-            mkp_ids0 = assign_keypoints(kpts0, cpdict[name0], conf['max_error'],
+            mkp_ids0 = assign_keypoints(kpts0, cpdict[name0], max_error0,
                                         update0, bindict[name0], scores,
-                                        conf['cell_size'])
-            mkp_ids1 = assign_keypoints(kpts1, cpdict[name1], max_error1,
+                                        cell_size0)
+            mkp_ids1 = assign_keypoints(kpts1, cpdict[name1], conf['max_error'],
                                         update1, bindict[name1], scores,
                                         conf['cell_size'])
 

hloc/matchers/loftr.py

@@ -23,9 +23,19 @@ def _init(self, conf):
         self.net = LoFTR_(pretrained=conf['weights'], config=cfg)
 
     def _forward(self, data):
+        # For consistency with hloc pairs, we refine kpts in image0!
+        rename = {
+            'keypoints0': 'keypoints1',
+            'keypoints1': 'keypoints0',
+            'image0': 'image1',
+            'image1': 'image0',
+            'mask0': 'mask1',
+            'mask1': 'mask0',
+        }
+        data_ = {rename[k]: v for k, v in data.items()}
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            pred = self.net(data)
+            pred = self.net(data_)
 
         scores = pred['confidence']
 
@@ -36,6 +46,8 @@ def _forward(self, data):
                 pred['keypoints0'][keep], pred['keypoints1'][keep]
             scores = scores[keep]
 
+        # Switch back indices
+        pred = {(rename[k] if k in rename else k): v for k, v in pred.items()}
         pred['scores'] = scores
         del pred['confidence']
         return pred