fix: twirl no match return None

lgrcia · Jul 21, 2023 · 12b9e04 · 12b9e04
1 parent e99c4b9
commit 12b9e04
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/prose/blocks/geometry.py b/prose/blocks/geometry.py
@@ -192,7 +192,7 @@ def run(self, image):
         else:
             image.discard = True
 
-    def solve(self, coords, tolerance=2):
+    def solve(self, coords, tolerance=2, refine=True):
         quads_image, asterisms_image = quads.hashes(coords)
         tree_image = cKDTree(quads_image)
         min_match = 0.7
@@ -216,17 +216,21 @@ def solve(self, coords, tolerance=2):
                     if match >= min_match * len(coords):
                         break
 
-        i, j = pairs[np.argmax(matches)]
-
-        if True:
-            M = get_transform_matrix(self.asterisms_ref[j], asterisms_image[i])
-            test = (M @ pad(self.ref).T)[0:2].T
-            s1, s2 = cross_match(coords, test, tolerance=tolerance, return_idxs=True).T
-            M = get_transform_matrix(self.ref[s2], coords[s1])
+        if len(matches) == 0:
+            return None
         else:
-            M = get_transform_matrix(self.asterisms_ref[j], asterisms_image[i])
+            i, j = pairs[np.argmax(matches)]
+            if refine:
+                M = get_transform_matrix(self.asterisms_ref[j], asterisms_image[i])
+                test = (M @ pad(self.ref).T)[0:2].T
+                s1, s2 = cross_match(
+                    coords, test, tolerance=tolerance, return_idxs=True
+                ).T
+                M = get_transform_matrix(self.ref[s2], coords[s1])
+            else:
+                M = get_transform_matrix(self.asterisms_ref[j], asterisms_image[i])
 
-        return M
+            return M
 
 
 # backward compatibility