update: lanet

hloc/extractors/lanet.py

@@ -34,10 +34,10 @@ class LANet(BaseModel):
 
         # Scores & Descriptors
         kpts_score = (
-            torch.cat([keypoints, scores], dim=1).view(3, -1).t().cpu().detach().numpy()
+            torch.cat([keypoints, scores], dim=1).view(3, -1).t()
         )
         descriptors = (
-            descriptors.view(256, Hc, Wc).view(256, -1).t().cpu().detach().numpy()
+            descriptors.view(256, Hc, Wc).view(256, -1).t()
         )
 
         # Filter based on confidence threshold
@@ -46,6 +46,11 @@ class LANet(BaseModel):
         keypoints = kpts_score[:, 1:]
         scores = kpts_score[:, 0]
 
+        idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+        keypoints = keypoints[idxs, :2]
+        descriptors = descriptors[idxs]
+        scores = scores[idxs]
+
         return {
             "keypoints": torch.from_numpy(keypoints)[None],
             "scores": torch.from_numpy(scores)[None],