fix: cache inference

.gitignore

@@ -18,3 +18,5 @@ gradio_cached_examples
 hloc/matchers/quadtree.py
 third_party/QuadTreeAttention
 desktop.ini
+experiments*
+datasets/wxbs_benchmark

common/utils.py

@@ -518,7 +518,7 @@ def run_matching(
     gr.Info(f"Matching images done using: {time.time()-t1:.3f}s")
     logger.info(f"Matching images done using: {time.time()-t1:.3f}s")
     t1 = time.time()
-    # plot images with keypoints
+    # plot images with keypoints\
     titles = [
         "Image 0 - Keypoints",
         "Image 1 - Keypoints",

common/viz.py

@@ -293,7 +293,7 @@ def draw_matches_core(
                 mkpts1,
                 color,
                 titles=titles,
-                # text=texts,
+                text=text,
                 path=path,
                 dpi=dpi,
                 pad=pad,
@@ -308,7 +308,7 @@ def draw_matches_core(
                 mkpts1,
                 color,
                 titles=titles,
-                # text=texts,
+                text=text,
                 pad=pad,
                 dpi=dpi,
             )

hloc/extractors/superpoint.py

@@ -44,4 +44,4 @@ class SuperPoint(BaseModel):
         self.net = superpoint.SuperPoint(conf)
 
     def _forward(self, data):
-        return self.net(data)
+        return self.net(data, self.conf)

third_party/SuperGluePretrainedNetwork/models/superpoint.py

@@ -83,9 +83,9 @@ def sample_descriptors(keypoints, descriptors, s: int = 8):
     """Interpolate descriptors at keypoint locations"""
     b, c, h, w = descriptors.shape
     keypoints = keypoints - s / 2 + 0.5
-    keypoints /= torch.tensor([(w * s - s / 2 - 0.5), (h * s - s / 2 - 0.5)],).to(
-        keypoints
-    )[None]
+    keypoints /= torch.tensor(
+        [(w * s - s / 2 - 0.5), (h * s - s / 2 - 0.5)],
+    ).to(keypoints)[None]
     keypoints = keypoints * 2 - 1  # normalize to (-1, 1)
     args = {"align_corners": True} if torch.__version__ >= "1.3" else {}
     descriptors = torch.nn.functional.grid_sample(
@@ -136,7 +136,11 @@ class SuperPoint(nn.Module):
 
         self.convDa = nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)
         self.convDb = nn.Conv2d(
-            c5, self.config["descriptor_dim"], kernel_size=1, stride=1, padding=0
+            c5,
+            self.config["descriptor_dim"],
+            kernel_size=1,
+            stride=1,
+            padding=0,
         )
 
         path = Path(__file__).parent / "weights/superpoint_v1.pth"
@@ -148,8 +152,12 @@ class SuperPoint(nn.Module):
 
         print("Loaded SuperPoint model")
 
-    def forward(self, data):
+    def forward(self, data, cfg={}):
         """Compute keypoints, scores, descriptors for image"""
+        self.config = {
+            **self.config,
+            **cfg,
+        }
         # Shared Encoder
         x = self.relu(self.conv1a(data["image"]))
         x = self.relu(self.conv1b(x))
@@ -182,7 +190,9 @@ class SuperPoint(nn.Module):
         keypoints, scores = list(
             zip(
                 *[
-                    remove_borders(k, s, self.config["remove_borders"], h * 8, w * 8)
+                    remove_borders(
+                        k, s, self.config["remove_borders"], h * 8, w * 8
+                    )
                     for k, s in zip(keypoints, scores)
                 ]
             )