feat: check if slate is valide

UCSD-E4E · Sep 15, 2024 · 4301b95 · 4301b95
1 parent 05e9945
commit 4301b95
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 26 deletions.
diff --git a/pyfishsensedev/library/homography/image_matcher.py b/pyfishsensedev/library/homography/image_matcher.py
@@ -67,8 +67,8 @@ def __call__(self, image1: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
             matcher = LightGlue(**self.matcher_conf).eval().to(device)
 
             # Extract features
-            if not self.com_license:
-                print("    WARNING: USING THE NON-COMMERCIAL VERSION OF SUPERPOINT!")
+            # if not self.com_license:
+            #     print("    WARNING: USING THE NON-COMMERCIAL VERSION OF SUPERPOINT!")
             if (
                 self.feats0 == None
             ):  # so we don't have to extract template keypoints every time

diff --git a/pyfishsensedev/plane_detector/slate_detector.py b/pyfishsensedev/plane_detector/slate_detector.py
@@ -1,4 +1,7 @@
+from typing import Tuple
+
 import numpy as np
+import torch
 
 from pyfishsensedev.image.pdf import Pdf
 from pyfishsensedev.library.homography.image_matcher import ImageMatcher
@@ -12,40 +15,50 @@ def __init__(self, image: np.ndarray[np.uint8], pdf: Pdf) -> None:
 
         self._pdf = pdf
         self._tensor_template = numpy_image_to_torch(self._pdf.image)
+        self._ran_template_matches = False
+        self._feats0_matches = None
+        self._feats1_matches = None
+
+    def _get_template_matches(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        if not self._ran_template_matches:
+            max_num_keypoints = 0
 
-    def _get_template_matches(self):
-        max_num_keypoints = 0
-        max_feats0_matches = None
-        max_feats1_matches = None
+            for scale in range(1, 20):
+                scale = float(scale) / 10.0
 
-        for scale in range(1, 20):
-            scale = float(scale) / 10.0
+                image_matcher = ImageMatcher(
+                    self._tensor_template,
+                    com_license=False,
+                    processing_conf={
+                        "preprocess": {"gamma": 2.0, "sharpness": None, "scale": scale},
+                        "matcher": {"filter_threshold": 0.5},
+                    },
+                )
 
-            image_matcher = ImageMatcher(
-                self._tensor_template,
-                processing_conf={
-                    "preprocess": {"gamma": 2.0, "sharpness": None, "scale": scale},
-                    "extractor": {"max_num_keypoints": None},
-                    "matcher": {"filter_threshold": 0.5},
-                },
-            )
+                feats0_matches, feats1_matches = image_matcher(
+                    numpy_image_to_torch(self.image)
+                )
 
-            feats0_matches, feats1_matches = image_matcher(
-                numpy_image_to_torch(self.image)
-            )
+                num_keypoints, _ = feats0_matches.shape
+                if num_keypoints > max_num_keypoints:
+                    max_num_keypoints = num_keypoints
 
-            num_keypoints, _ = feats0_matches.shape
-            if num_keypoints > max_num_keypoints:
-                max_num_keypoints = num_keypoints
+                    self._max_feats0_matches = feats0_matches
+                    self._max_feats1_matches = feats1_matches
 
-                max_feats0_matches = feats0_matches
-                max_feats1_matches = feats1_matches
+            self._ran_template_matches = True
 
-        if max_feats0_matches is not None:
-            print(f"{max_feats0_matches.size()}, {max_feats1_matches.size()}")
+        return self._max_feats0_matches, self._feats1_matches
 
     def _get_points_image_space(self):
         pass
 
     def _get_points_body_space(self):
         pass
+
+    def is_valid(self):
+        feats0_matches, _ = self._get_template_matches()
+
+        num_matches, _ = feats0_matches.shape
+
+        return num_matches > 10