Fixed fisheye

albumentations/augmentations/geometric/functional.py

@@ -3357,21 +3357,22 @@ def tps_transform(
 def get_camera_matrix_distortion_maps(
     image_shape: tuple[int, int],
     k: float,
+    center_xy: tuple[float, float],
 ) -> tuple[np.ndarray, np.ndarray]:
     """Generate distortion maps using camera matrix model.
 
     Args:
         image_shape: Image shape
         k: Distortion coefficient
-
+        center_xy: Center of distortion
     Returns:
         tuple of:
         - map_x: Horizontal displacement map
         - map_y: Vertical displacement map
     """
     height, width = image_shape[:2]
     camera_matrix = np.array(
-        [[width, 0, 0], [0, height, 0], [0, 0, 1]],
+        [[width, 0, center_xy[0]], [0, height, center_xy[1]], [0, 0, 1]],
         dtype=np.float32,
     )
     distortion = np.array([k, k, 0, 0, 0], dtype=np.float32)
@@ -3388,31 +3389,42 @@ def get_camera_matrix_distortion_maps(
 def get_fisheye_distortion_maps(
     image_shape: tuple[int, int],
     k: float,
+    center_xy: tuple[float, float],
 ) -> tuple[np.ndarray, np.ndarray]:
     """Generate distortion maps using fisheye model.
 
     Args:
         image_shape: Image shape
         k: Distortion coefficient
-
+        center_xy: Center of distortion
     Returns:
         tuple of:
         - map_x: Horizontal displacement map
         - map_y: Vertical displacement map
     """
     height, width = image_shape[:2]
+
+    center_x, center_y = center_xy
+
     # Create coordinate grid
     y, x = np.mgrid[:height, :width].astype(np.float32)
 
+    x = x - center_x
+    y = y - center_y
+
     # Calculate polar coordinates
     r = np.sqrt(x * x + y * y)
     theta = np.arctan2(y, x)
 
-    # Apply fisheye distortion
-    r_dist = r * (1 + k * r * r)
+    # Normalize radius by the maximum possible radius to keep distortion in check
+    max_radius = math.sqrt(max(center_x, width - center_x) ** 2 + max(center_y, height - center_y) ** 2)
+    r_norm = r / max_radius
+
+    # Apply fisheye distortion to normalized radius
+    r_dist = r * (1 + k * r_norm * r_norm)
 
     # Convert back to cartesian coordinates
-    map_x = r_dist * np.cos(theta)
-    map_y = r_dist * np.sin(theta)
+    map_x = r_dist * np.cos(theta) + center_x
+    map_y = r_dist * np.sin(theta) + center_y
 
     return map_x, map_y

albumentations/augmentations/geometric/transforms.py

@@ -1622,17 +1622,20 @@ def get_params_dependent_on_data(
         k = self.py_random.uniform(*self.distort_limit)
 
         # Calculate center shift
+        center_xy = fgeometric.center(image_shape)
 
         # Get distortion maps based on mode
         if self.mode == "camera":
             map_x, map_y = fgeometric.get_camera_matrix_distortion_maps(
                 image_shape,
                 k,
+                center_xy,
             )
         else:  # fisheye
             map_x, map_y = fgeometric.get_fisheye_distortion_maps(
                 image_shape,
                 k,
+                center_xy,
             )
 
         return {"map_x": map_x, "map_y": map_y}