diff --git a/albumentations/augmentations/geometric/functional.py b/albumentations/augmentations/geometric/functional.py index dcdf1b0f8..797f54dfe 100644 --- a/albumentations/augmentations/geometric/functional.py +++ b/albumentations/augmentations/geometric/functional.py @@ -3357,13 +3357,14 @@ 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 @@ -3371,7 +3372,7 @@ def get_camera_matrix_distortion_maps( """ 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 diff --git a/albumentations/augmentations/geometric/transforms.py b/albumentations/augmentations/geometric/transforms.py index d0f0ed1c5..1436840b5 100644 --- a/albumentations/augmentations/geometric/transforms.py +++ b/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}