add: extract

Vincentqyw · Sep 25, 2024 · 0f1f0fd · 0f1f0fd
1 parent 0e244d3
commit 0f1f0fd
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/ui/api.py b/ui/api.py
@@ -147,6 +147,25 @@ def _forward(self, img0, img1):
             pred = match_features.match_images(self.matcher, pred0, pred1)
         return pred
 
+    @torch.inference_mode()
+    def extract(
+        self,
+        img0: np.ndarray,
+    ) -> Dict[str, np.ndarray]:
+        """Extract features from a single image.
+
+        Args:
+            img0 (np.ndarray): image
+
+        Returns:
+            Dict[str, np.ndarray]: feature dict
+        """
+
+        pred = extract_features.extract(
+            self.extractor, img0, self.extract_conf["preprocessing"]
+        )
+        return pred
+
     @torch.inference_mode()
     def forward(
         self,