Merge pull request #90 from UCL/89-optional-undistort

89 optional undistort
SciKit-Surgery · Feb 2, 2021 · 74f760c · 74f760c
2 parents 135f9d3 + 8367729
commit 74f760c
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 23 deletions.
diff --git a/sksurgeryimage/calibration/aruco_point_detector.py b/sksurgeryimage/calibration/aruco_point_detector.py
@@ -62,7 +62,7 @@ def __init__(self,
         self.parameters = parameters
         self.model = model
 
-    def _internal_get_points(self, image):
+    def _internal_get_points(self, image, is_distorted=True):
         """
         Extracts points using OpenCV's ArUco implementation.
 

diff --git a/sksurgeryimage/calibration/charuco_plus_chessboard_point_detector.py b/sksurgeryimage/calibration/charuco_plus_chessboard_point_detector.py
@@ -135,7 +135,7 @@ def __init__(self,
         if self.model_points.shape[0] == 0:
             raise ValueError("No reference model points were found")
 
-    def _internal_get_points(self, image):
+    def _internal_get_points(self, image, is_distorted=True):
         """
         Extracts points using scikit-surgeryimage ChArUcoPointDetector
         and ChessboardPointDetector classes.

diff --git a/sksurgeryimage/calibration/charuco_point_detector.py b/sksurgeryimage/calibration/charuco_point_detector.py
@@ -63,7 +63,7 @@ def __init__(self, dictionary,
                                        * self.size[0]
             self.object_points[i][2] = 0
 
-    def _internal_get_points(self, image):
+    def _internal_get_points(self, image, is_distorted=True):
         """
         Extracts points using OpenCV's ChArUco implementation.
 

diff --git a/sksurgeryimage/calibration/chessboard_point_detector.py b/sksurgeryimage/calibration/chessboard_point_detector.py
@@ -42,7 +42,7 @@ def __init__(self, number_of_corners, square_size_in_mm, scale=(1, 1)):
             self.object_points[i][2] = 0
             self.ids[i][0] = i
 
-    def _internal_get_points(self, image):
+    def _internal_get_points(self, image, is_distorted=True):
         """
         Extracts points using OpenCV's chessboard implementation.
 

diff --git a/sksurgeryimage/calibration/dotty_grid_point_detector.py b/sksurgeryimage/calibration/dotty_grid_point_detector.py
@@ -114,11 +114,13 @@ def __init__(self,
         self.min_area = min_area
         self.max_area = max_area
 
-    def _internal_get_points(self, image):
+    def _internal_get_points(self, image, is_distorted=True):
         """
         Extracts points.
 
         :param image: numpy 2D grey scale image.
+        :param is_distorted: False if the input image has already been \
+             undistorted.
         :return: ids, object_points, image_points as Nx[1,3,2] ndarrays
         """
 
@@ -148,21 +150,29 @@ def _internal_get_points(self, image):
         detector = cv2.SimpleBlobDetector_create(params)
         keypoints = detector.detect(thresholded)
 
-        # Also detect them in an undistorted image.
-        undistorted_image = cv2.undistort(smoothed,
-                                          self.intrinsics,
-                                          self.distortion_coefficients
-                                          )
-
-        undistorted_thresholded = \
-            cv2.adaptiveThreshold(undistorted_image,
-                                  255,
-                                  cv2.ADAPTIVE_THRESH_MEAN_C,
-                                  cv2.THRESH_BINARY,
-                                  self.threshold_window_size,
-                                  self.threshold_offset)
-
-        undistorted_keypoints = detector.detect(undistorted_thresholded)
+        # If input image is distorted, undistort and also detect points
+        # in undistorted image.
+        if is_distorted:
+            undistorted_image = cv2.undistort(smoothed,
+                                              self.intrinsics,
+                                              self.distortion_coefficients
+                                              )
+
+
+            undistorted_thresholded = \
+                cv2.adaptiveThreshold(undistorted_image,
+                                      255,
+                                      cv2.ADAPTIVE_THRESH_MEAN_C,
+                                      cv2.THRESH_BINARY,
+                                      self.threshold_window_size,
+                                      self.threshold_offset)
+
+            undistorted_keypoints = detector.detect(undistorted_thresholded)
+
+        else:
+            undistorted_image = image
+            undistorted_thresholded = thresholded
+            undistorted_keypoints = keypoints
 
         # Note that keypoints and undistorted_keypoints
         # can be of different length

diff --git a/sksurgeryimage/calibration/point_detector.py b/sksurgeryimage/calibration/point_detector.py
@@ -23,11 +23,13 @@ def __init__(self, scale=(1, 1)):
         self.scale = scale
         self.scale_x, self.scale_y = scale
 
-    def get_points(self, image):
+    def get_points(self, image, is_distorted=True):
         """
         Client's call this method to extract points from an image.
 
         :param image: numpy 2D RGB/grayscale image.
+        :param is_distorted: False if the input image has already been \
+             undistorted.
         :return: ids, object_points, image_points as Nx[1,3,2] ndarrays
         """
 
@@ -51,19 +53,22 @@ def get_points(self, image):
         else:
             resized = grey
 
-        ids, object_points, image_points = self._internal_get_points(resized)
+        ids, object_points, image_points = \
+            self._internal_get_points(resized, is_distorted=is_distorted)
 
         if is_resized:
             image_points[:, 0] /= self.scale_x
             image_points[:, 1] /= self.scale_y
 
         return ids, object_points, image_points
 
-    def _internal_get_points(self, image):
+    def _internal_get_points(self, image, is_distorted=True):
         """
         Derived classes override this one.
 
         :param image: numpy 2D grey scale image.
+        :param is_distorted: False if the input image has already been \
+             undistorted.
         :return: ids, object_points, image_points as Nx[1,3,2] ndarrays
         """
         raise NotImplementedError()