refactor: adding enum for rectification options

tier4 · Jan 17, 2025 · c7026ec · c7026ec
1 parent 00c929a
commit c7026ec
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/...a_calibrator/intrinsic_camera_calibrator/intrinsic_camera_calibrator/camera_calibrator.py b/...a_calibrator/intrinsic_camera_calibrator/intrinsic_camera_calibrator/camera_calibrator.py
@@ -62,6 +62,7 @@
 from intrinsic_camera_calibrator.parameter import ParameterizedClass
 from intrinsic_camera_calibrator.types import ImageViewMode
 from intrinsic_camera_calibrator.types import OperationMode
+from intrinsic_camera_calibrator.types import RectifyMode
 from intrinsic_camera_calibrator.utils import save_intrinsics
 from intrinsic_camera_calibrator.utils import set_logger_severity
 from intrinsic_camera_calibrator.views.data_collector_view import DataCollectorView
@@ -170,8 +171,8 @@ def make_mode_group(self):
 
         self.rectify_label = QLabel("Rectify option:")
         self.rectify_type_combobox = QComboBox()
-        self.rectify_type_combobox.addItem("OpenCV", 0)
-        self.rectify_type_combobox.addItem("Fixed aspect ratio", 1)
+        self.rectify_type_combobox.addItem(RectifyMode.OPENCV.value, RectifyMode.OPENCV)
+        self.rectify_type_combobox.addItem(RectifyMode.FIXED_ASPECT_RATIO.value, RectifyMode.FIXED_ASPECT_RATIO)
         self.rectify_type_combobox.setEnabled(False)
 
         def pause_callback():

diff --git a/...tor/intrinsic_camera_calibrator/intrinsic_camera_calibrator/camera_models/camera_model.py b/...tor/intrinsic_camera_calibrator/intrinsic_camera_calibrator/camera_models/camera_model.py
@@ -19,6 +19,7 @@
 from typing import List
 from typing import Optional
 from typing import Tuple
+from intrinsic_camera_calibrator.types import RectifyMode
 
 import cv2
 import numpy as np
@@ -211,7 +212,7 @@ def update_config(self, **kwargs):
 
     def get_undistorted_camera_model(self, alpha: float, rectify_option):
         """Compute the undistorted version of the camera model."""
-        if rectify_option == 1:
+        if rectify_option == RectifyMode.FIXED_ASPECT_RATIO:
 
             def get_rectangles(camera_matrix, dist_coeffs, img_size, new_camera_matrix=None):
                 N = 101
@@ -315,7 +316,7 @@ def roi_to_intrinsics(roi):
                 height=new_image_height,
                 width=new_image_width,
             )
-        elif rectify_option == 0:
+        elif rectify_option == RectifyMode.OPENCV:
             undistorted_k, _ = cv2.getOptimalNewCameraMatrix(
                 self.k, self.d, (self.width, self.height), alpha
             )

diff --git a/...rinsic_camera_calibrator/intrinsic_camera_calibrator/intrinsic_camera_calibrator/types.py b/...rinsic_camera_calibrator/intrinsic_camera_calibrator/intrinsic_camera_calibrator/types.py
@@ -42,3 +42,9 @@ class CollectionStatus(Enum):
     REJECTED = 2
     REDUNDANT = 3
     ACCEPTED = 4
+
+class RectifyMode(Enum):
+    """Methods to rectify the image."""
+
+    OPENCV = "OpenCV"
+    FIXED_ASPECT_RATIO = "Fixed aspect ratio"