From 0f1f0fdf06b1335ddfe1d8003233411deff4d8fa Mon Sep 17 00:00:00 2001
From: Realcat <alpharealcat@gmail.com>
Date: Wed, 25 Sep 2024 01:48:29 +0000
Subject: [PATCH] add: extract

---
 ui/api.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/ui/api.py b/ui/api.py
index bcbd697..03f68ae 100644
--- 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,