nrtk-albumentations 2.1.0__py3-none-any.whl

albumentations/augmentations/geometric/functional.py

@@ -0,0 +1,4151 @@
+"""Functional implementations of geometric image transformations.
+
+This module provides low-level functions for geometric operations such as rotation,
+resizing, flipping, perspective transforms, and affine transformations on images,
+bounding boxes and keypoints.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import defaultdict
+from collections.abc import Mapping, Sequence
+from typing import Any, Literal, cast
+from warnings import warn
+
+import cv2
+import numpy as np
+from albucore import (
+    get_num_channels,
+    hflip,
+    maybe_process_in_chunks,
+    preserve_channel_dim,
+    vflip,
+)
+
+from albumentations.augmentations.utils import angle_2pi_range, handle_empty_array
+from albumentations.core.bbox_utils import (
+    bboxes_from_masks,
+    bboxes_to_mask,
+    denormalize_bboxes,
+    mask_to_bboxes,
+    masks_from_bboxes,
+    normalize_bboxes,
+)
+from albumentations.core.type_definitions import (
+    NUM_BBOXES_COLUMNS_IN_ALBUMENTATIONS,
+    NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS,
+    NUM_MULTI_CHANNEL_DIMENSIONS,
+    REFLECT_BORDER_MODES,
+)
+
+PAIR = 2
+
+ROT90_180_FACTOR = 2
+ROT90_270_FACTOR = 3
+
+
+@handle_empty_array("bboxes")
+def bboxes_rot90(bboxes: np.ndarray, factor: int) -> np.ndarray:
+    """Rotates bounding boxes by 90 degrees CCW (see np.rot90)
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (num_boxes, 4+)
+        factor (int): Number of 90-degree rotations (1, 2, or 3)
+
+    Returns:
+        np.ndarray: Rotated bounding boxes
+
+    """
+    if factor == 0:
+        return bboxes
+
+    rotated_bboxes = bboxes.copy()
+    x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]
+
+    if factor == 1:
+        rotated_bboxes[:, 0] = y_min
+        rotated_bboxes[:, 1] = 1 - x_max
+        rotated_bboxes[:, 2] = y_max
+        rotated_bboxes[:, 3] = 1 - x_min
+    elif factor == ROT90_180_FACTOR:
+        rotated_bboxes[:, 0] = 1 - x_max
+        rotated_bboxes[:, 1] = 1 - y_max
+        rotated_bboxes[:, 2] = 1 - x_min
+        rotated_bboxes[:, 3] = 1 - y_min
+    elif factor == ROT90_270_FACTOR:
+        rotated_bboxes[:, 0] = 1 - y_max
+        rotated_bboxes[:, 1] = x_min
+        rotated_bboxes[:, 2] = 1 - y_min
+        rotated_bboxes[:, 3] = x_max
+
+    return rotated_bboxes
+
+
+@handle_empty_array("bboxes")
+def bboxes_d4(
+    bboxes: np.ndarray,
+    group_member: Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"],
+) -> np.ndarray:
+    """Applies a `D_4` symmetry group transformation to a bounding box.
+
+    The function transforms a bounding box according to the specified group member from the `D_4` group.
+    These transformations include rotations and reflections, specified to work on an image's bounding box given
+    its dimensions.
+
+    Args:
+        bboxes (np.ndarray): A numpy array of bounding boxes with shape (num_bboxes, 4+).
+        Each row represents a bounding box (x_min, y_min, x_max, y_max, ...).
+        group_member (Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"]): A string identifier for the
+            `D_4` group transformation to apply.
+
+    Returns:
+        BoxInternalType: The transformed bounding box.
+
+    Raises:
+        ValueError: If an invalid group member is specified.
+
+    """
+    transformations = {
+        "e": lambda x: x,  # Identity transformation
+        "r90": lambda x: bboxes_rot90(x, 1),  # Rotate 90 degrees
+        "r180": lambda x: bboxes_rot90(x, 2),  # Rotate 180 degrees
+        "r270": lambda x: bboxes_rot90(x, 3),  # Rotate 270 degrees
+        "v": lambda x: bboxes_vflip(x),  # Vertical flip
+        "hvt": lambda x: bboxes_transpose(
+            bboxes_rot90(x, 2),
+        ),  # Reflect over anti-diagonal
+        "h": lambda x: bboxes_hflip(x),  # Horizontal flip
+        "t": lambda x: bboxes_transpose(x),  # Transpose (reflect over main diagonal)
+    }
+
+    # Execute the appropriate transformation
+    if group_member in transformations:
+        return transformations[group_member](bboxes)
+
+    raise ValueError(f"Invalid group member: {group_member}")
+
+
+@handle_empty_array("keypoints")
+@angle_2pi_range
+def keypoints_rot90(
+    keypoints: np.ndarray,
+    factor: Literal[0, 1, 2, 3],
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Rotate keypoints by 90 degrees counter-clockwise (CCW) a specified number of times.
+
+    Args:
+        keypoints (np.ndarray): An array of keypoints with shape (N, 4+) in the format (x, y, angle, scale, ...).
+        factor (int): The number of 90 degree CCW rotations to apply. Must be in the range [0, 3].
+        image_shape (tuple[int, int]): The shape of the image (height, width).
+
+    Returns:
+        np.ndarray: The rotated keypoints with the same shape as the input.
+
+    """
+    if factor == 0:
+        return keypoints
+
+    height, width = image_shape[:2]
+    rotated_keypoints = keypoints.copy().astype(np.float32)
+
+    x, y, angle = keypoints[:, 0], keypoints[:, 1], keypoints[:, 3]
+
+    if factor == 1:
+        rotated_keypoints[:, 0] = y
+        rotated_keypoints[:, 1] = width - 1 - x
+        rotated_keypoints[:, 3] = angle - np.pi / 2
+    elif factor == ROT90_180_FACTOR:
+        rotated_keypoints[:, 0] = width - 1 - x
+        rotated_keypoints[:, 1] = height - 1 - y
+        rotated_keypoints[:, 3] = angle - np.pi
+    elif factor == ROT90_270_FACTOR:
+        rotated_keypoints[:, 0] = height - 1 - y
+        rotated_keypoints[:, 1] = x
+        rotated_keypoints[:, 3] = angle + np.pi / 2
+
+    return rotated_keypoints
+
+
+@handle_empty_array("keypoints")
+def keypoints_d4(
+    keypoints: np.ndarray,
+    group_member: Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"],
+    image_shape: tuple[int, int],
+    **params: Any,
+) -> np.ndarray:
+    """Applies a `D_4` symmetry group transformation to a keypoint.
+
+    This function adjusts a keypoint's coordinates according to the specified `D_4` group transformation,
+    which includes rotations and reflections suitable for image processing tasks. These transformations account
+    for the dimensions of the image to ensure the keypoint remains within its boundaries.
+
+    Args:
+        keypoints (np.ndarray): An array of keypoints with shape (N, 4+) in the format (x, y, angle, scale, ...).
+        group_member (Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"]): A string identifier for
+            the `D_4` group transformation to apply.
+            Valid values are 'e', 'r90', 'r180', 'r270', 'v', 'hv', 'h', 't'.
+        image_shape (tuple[int, int]): The shape of the image.
+        params (Any): Not used.
+
+    Returns:
+        KeypointInternalType: The transformed keypoint.
+
+    Raises:
+        ValueError: If an invalid group member is specified, indicating that the specified transformation
+            does not exist.
+
+    """
+    rows, cols = image_shape[:2]
+    transformations = {
+        "e": lambda x: x,  # Identity transformation
+        "r90": lambda x: keypoints_rot90(x, 1, image_shape),  # Rotate 90 degrees
+        "r180": lambda x: keypoints_rot90(x, 2, image_shape),  # Rotate 180 degrees
+        "r270": lambda x: keypoints_rot90(x, 3, image_shape),  # Rotate 270 degrees
+        "v": lambda x: keypoints_vflip(x, rows),  # Vertical flip
+        "hvt": lambda x: keypoints_transpose(
+            keypoints_rot90(x, 2, image_shape),
+        ),  # Reflect over anti diagonal
+        "h": lambda x: keypoints_hflip(x, cols),  # Horizontal flip
+        "t": lambda x: keypoints_transpose(x),  # Transpose (reflect over main diagonal)
+    }
+    # Execute the appropriate transformation
+    if group_member in transformations:
+        return transformations[group_member](keypoints)
+
+    raise ValueError(f"Invalid group member: {group_member}")
+
+
+@preserve_channel_dim
+def resize(
+    img: np.ndarray,
+    target_shape: tuple[int, int],
+    interpolation: int,
+) -> np.ndarray:
+    """Resize an image to the specified dimensions.
+
+    This function resizes an input image to the target shape using the specified
+    interpolation method. If the image is already the target size, it is returned unchanged.
+
+    Args:
+        img (np.ndarray): Input image to resize.
+        target_shape (tuple[int, int]): Target (height, width) dimensions.
+        interpolation (int): Interpolation method to use (cv2 interpolation flag).
+            Examples: cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_NEAREST, etc.
+
+    Returns:
+        np.ndarray: Resized image with shape target_shape + original channel dimensions.
+
+    """
+    if target_shape == img.shape[:2]:
+        return img
+
+    height, width = target_shape[:2]
+    resize_fn = maybe_process_in_chunks(
+        cv2.resize,
+        dsize=(width, height),
+        interpolation=interpolation,
+    )
+    return resize_fn(img)
+
+
+@preserve_channel_dim
+def scale(img: np.ndarray, scale: float, interpolation: int) -> np.ndarray:
+    """Scale an image by a factor while preserving aspect ratio.
+
+    This function scales both height and width dimensions of the image by the same factor.
+
+    Args:
+        img (np.ndarray): Input image to scale.
+        scale (float): Scale factor. Values > 1 will enlarge the image, values < 1 will shrink it.
+        interpolation (int): Interpolation method to use (cv2 interpolation flag).
+
+    Returns:
+        np.ndarray: Scaled image.
+
+    """
+    height, width = img.shape[:2]
+    new_size = int(height * scale), int(width * scale)
+    return resize(img, new_size, interpolation)
+
+
+@handle_empty_array("keypoints")
+def keypoints_scale(
+    keypoints: np.ndarray,
+    scale_x: float,
+    scale_y: float,
+) -> np.ndarray:
+    """Scale keypoints by given factors.
+
+    Args:
+        keypoints (np.ndarray): Array of keypoints with shape (num_keypoints, 2+)
+        scale_x (float): Scale factor for x coordinates
+        scale_y (float): Scale factor for y coordinates
+
+    Returns:
+        np.ndarray: Scaled keypoints
+
+    """
+    # Extract x, y, z, angle, and scale
+    x, y, z, angle, scale = (
+        keypoints[:, 0],
+        keypoints[:, 1],
+        keypoints[:, 2],
+        keypoints[:, 3],
+        keypoints[:, 4],
+    )
+
+    # Scale x and y
+    x_scaled = x * scale_x
+    y_scaled = y * scale_y
+
+    # Scale the keypoint scale by the maximum of scale_x and scale_y
+    scale_scaled = scale * max(scale_x, scale_y)
+
+    # Create the output array
+    scaled_keypoints = np.column_stack([x_scaled, y_scaled, z, angle, scale_scaled])
+
+    # If there are additional columns, preserve them
+    if keypoints.shape[1] > NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS:
+        return np.column_stack(
+            [scaled_keypoints, keypoints[:, NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS:]],
+        )
+
+    return scaled_keypoints
+
+
+@preserve_channel_dim
+def perspective(
+    img: np.ndarray,
+    matrix: np.ndarray,
+    max_width: int,
+    max_height: int,
+    border_val: float | list[float] | np.ndarray,
+    border_mode: int,
+    keep_size: bool,
+    interpolation: int,
+) -> np.ndarray:
+    """Apply perspective transformation to an image.
+
+    This function warps an image according to a perspective transformation matrix.
+    It can either maintain the original dimensions or use the specified max dimensions.
+
+    Args:
+        img (np.ndarray): Input image to transform.
+        matrix (np.ndarray): 3x3 perspective transformation matrix.
+        max_width (int): Maximum width of the output image if keep_size is False.
+        max_height (int): Maximum height of the output image if keep_size is False.
+        border_val (float | list[float] | np.ndarray): Border value(s) to fill areas outside the transformed image.
+        border_mode (int): OpenCV border mode (e.g., cv2.BORDER_CONSTANT, cv2.BORDER_REFLECT).
+        keep_size (bool): If True, maintain the original image dimensions.
+        interpolation (int): Interpolation method for resampling (cv2 interpolation flag).
+
+    Returns:
+        np.ndarray: Perspective-transformed image.
+
+    """
+    if not keep_size:
+        perspective_func = maybe_process_in_chunks(
+            cv2.warpPerspective,
+            M=matrix,
+            dsize=(max_width, max_height),
+            borderMode=border_mode,
+            borderValue=border_val,
+            flags=interpolation,
+        )
+    else:
+        height, width = img.shape[:2]
+
+        scale_x = width / max_width
+        scale_y = height / max_height
+        scale_matrix = np.array([[scale_x, 0, 0], [0, scale_y, 0], [0, 0, 1]])
+        adjusted_matrix = np.dot(scale_matrix, matrix)
+
+        perspective_func = maybe_process_in_chunks(
+            cv2.warpPerspective,
+            M=adjusted_matrix,
+            dsize=(width, height),
+            borderMode=border_mode,
+            borderValue=border_val,
+            flags=interpolation,
+        )
+
+    return perspective_func(img)
+
+
+@handle_empty_array("bboxes")
+def perspective_bboxes(
+    bboxes: np.ndarray,
+    image_shape: tuple[int, int],
+    matrix: np.ndarray,
+    max_width: int,
+    max_height: int,
+    keep_size: bool,
+) -> np.ndarray:
+    """Applies perspective transformation to bounding boxes.
+
+    This function transforms bounding boxes using the given perspective transformation matrix.
+    It handles bounding boxes with additional attributes beyond the standard coordinates.
+
+    Args:
+        bboxes (np.ndarray): An array of bounding boxes with shape (num_bboxes, 4+).
+                             Each row represents a bounding box (x_min, y_min, x_max, y_max, ...).
+                             Additional columns beyond the first 4 are preserved unchanged.
+        image_shape (tuple[int, int]): The shape of the image (height, width).
+        matrix (np.ndarray): The perspective transformation matrix.
+        max_width (int): The maximum width of the output image.
+        max_height (int): The maximum height of the output image.
+        keep_size (bool): If True, maintains the original image size after transformation.
+
+    Returns:
+        np.ndarray: An array of transformed bounding boxes with the same shape as input.
+                    The first 4 columns contain the transformed coordinates, and any
+                    additional columns are preserved from the input.
+
+    Note:
+        - This function modifies only the coordinate columns (first 4) of the input bounding boxes.
+        - Any additional attributes (columns beyond the first 4) are kept unchanged.
+        - The function handles denormalization and renormalization of coordinates internally.
+
+    Examples:
+        >>> bboxes = np.array([[0.1, 0.1, 0.3, 0.3, 1], [0.5, 0.5, 0.8, 0.8, 2]])
+        >>> image_shape = (100, 100)
+        >>> matrix = np.array([[1.5, 0.2, -20], [-0.1, 1.3, -10], [0.002, 0.001, 1]])
+        >>> transformed_bboxes = perspective_bboxes(bboxes, image_shape, matrix, 150, 150, False)
+
+    """
+    height, width = image_shape[:2]
+    transformed_bboxes = bboxes.copy()
+    denormalized_coords = denormalize_bboxes(bboxes[:, :4], image_shape)
+
+    x_min, y_min, x_max, y_max = denormalized_coords.T
+    points = np.array(
+        [[x_min, y_min], [x_max, y_min], [x_max, y_max], [x_min, y_max]],
+    ).transpose(2, 0, 1)
+    points_reshaped = points.reshape(-1, 1, 2)
+
+    transformed_points = cv2.perspectiveTransform(
+        points_reshaped.astype(np.float32),
+        matrix,
+    )
+    transformed_points = transformed_points.reshape(-1, 4, 2)
+
+    new_coords = np.array(
+        [[np.min(box[:, 0]), np.min(box[:, 1]), np.max(box[:, 0]), np.max(box[:, 1])] for box in transformed_points],
+    )
+
+    if keep_size:
+        scale_x, scale_y = width / max_width, height / max_height
+        new_coords[:, [0, 2]] *= scale_x
+        new_coords[:, [1, 3]] *= scale_y
+        output_shape = image_shape
+    else:
+        output_shape = (max_height, max_width)
+
+    normalized_coords = normalize_bboxes(new_coords, output_shape)
+    transformed_bboxes[:, :4] = normalized_coords
+
+    return transformed_bboxes
+
+
+def rotation2d_matrix_to_euler_angles(matrix: np.ndarray, y_up: bool) -> float:
+    """Args:
+    matrix (np.ndarray): Rotation matrix
+    y_up (bool): is Y axis looks up or down
+
+    """
+    if y_up:
+        return np.arctan2(matrix[1, 0], matrix[0, 0])
+    return np.arctan2(-matrix[1, 0], matrix[0, 0])
+
+
+@handle_empty_array("keypoints")
+@angle_2pi_range
+def perspective_keypoints(
+    keypoints: np.ndarray,
+    image_shape: tuple[int, int],
+    matrix: np.ndarray,
+    max_width: int,
+    max_height: int,
+    keep_size: bool,
+) -> np.ndarray:
+    """Apply perspective transformation to keypoints.
+
+    Args:
+        keypoints (np.ndarray): Array of shape (N, 5+) in format [x, y, z, angle, scale, ...]
+        image_shape (tuple[int, int]): Original image shape (height, width)
+        matrix (np.ndarray): 3x3 perspective transformation matrix
+        max_width (int): Maximum width after transformation
+        max_height (int): Maximum height after transformation
+        keep_size (bool): Whether to keep original size
+
+    Returns:
+        np.ndarray: Transformed keypoints array with same shape as input
+
+    """
+    keypoints = keypoints.copy().astype(np.float32)
+
+    height, width = image_shape[:2]
+
+    x, y, z, angle, scale = (
+        keypoints[:, 0],
+        keypoints[:, 1],
+        keypoints[:, 2],
+        keypoints[:, 3],
+        keypoints[:, 4],
+    )
+
+    # Reshape keypoints for perspective transform
+    keypoint_vector = np.column_stack((x, y)).astype(np.float32).reshape(-1, 1, 2)
+
+    # Apply perspective transform
+    transformed_points = cv2.perspectiveTransform(keypoint_vector, matrix).squeeze()
+
+    # Unsqueeze if we have a single keypoint
+    if transformed_points.ndim == 1:
+        transformed_points = transformed_points[np.newaxis, :]
+
+    x, y = transformed_points[:, 0], transformed_points[:, 1]
+
+    # Update angles
+    angle += rotation2d_matrix_to_euler_angles(matrix[:2, :2], y_up=True)
+
+    # Calculate scale factors
+    scale_x = np.sign(matrix[0, 0]) * np.sqrt(matrix[0, 0] ** 2 + matrix[0, 1] ** 2)
+    scale_y = np.sign(matrix[1, 1]) * np.sqrt(matrix[1, 0] ** 2 + matrix[1, 1] ** 2)
+    scale *= max(scale_x, scale_y)
+
+    if keep_size:
+        scale_x = width / max_width
+        scale_y = height / max_height
+        x *= scale_x
+        y *= scale_y
+        scale *= max(scale_x, scale_y)
+
+    # Create the output array with unchanged z coordinate
+    transformed_keypoints = np.column_stack([x, y, z, angle, scale])
+
+    # If there are additional columns, preserve them
+    if keypoints.shape[1] > NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS:
+        return np.column_stack(
+            [
+                transformed_keypoints,
+                keypoints[:, NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS:],
+            ],
+        )
+
+    return transformed_keypoints
+
+
+def is_identity_matrix(matrix: np.ndarray) -> bool:
+    """Check if the given matrix is an identity matrix.
+
+    Args:
+        matrix (np.ndarray): A 3x3 affine transformation matrix.
+
+    Returns:
+        bool: True if the matrix is an identity matrix, False otherwise.
+
+    """
+    return np.allclose(matrix, np.eye(3, dtype=matrix.dtype))
+
+
+def warp_affine_with_value_extension(
+    image: np.ndarray,
+    matrix: np.ndarray,
+    dsize: tuple[int, int],
+    flags: int,
+    border_mode: int,
+    border_value: tuple[float, ...] | float,
+) -> np.ndarray:
+    """Warp affine with value extension.
+
+    This function warps an image with a given affine transformation matrix.
+    It also extends the value to a sequence of floats.
+
+    Args:
+        image (np.ndarray): The image to warp.
+        matrix (np.ndarray): The affine transformation matrix.
+        dsize (tuple[int, int]): The size of the output image.
+        flags (int): The flags for the warp.
+        border_mode (int): The border mode to use.
+        border_value (tuple[float, ...] | float): The value to pad the image with.
+
+    Returns:
+        np.ndarray: The warped image.
+
+    """
+    num_channels = get_num_channels(image)
+    extended_value = extend_value(border_value, num_channels)
+
+    return cv2.warpAffine(
+        image,
+        matrix,
+        dsize,
+        flags=flags,
+        borderMode=border_mode,
+        borderValue=extended_value,
+    )
+
+
+@preserve_channel_dim
+def warp_affine(
+    image: np.ndarray,
+    matrix: np.ndarray,
+    interpolation: int,
+    fill: tuple[float, ...] | float,
+    border_mode: int,
+    output_shape: tuple[int, int],
+) -> np.ndarray:
+    """Apply an affine transformation to an image.
+
+    This function transforms an image using the specified affine transformation matrix.
+    If the transformation matrix is an identity matrix, the original image is returned.
+
+    Args:
+        image (np.ndarray): Input image to transform.
+        matrix (np.ndarray): 2x3 or 3x3 affine transformation matrix.
+        interpolation (int): Interpolation method for resampling.
+        fill (tuple[float, ...] | float): Border value(s) to fill areas outside the transformed image.
+        border_mode (int): OpenCV border mode for handling pixels outside the image boundaries.
+        output_shape (tuple[int, int]): Shape (height, width) of the output image.
+
+    Returns:
+        np.ndarray: Affine-transformed image with dimensions specified by output_shape.
+
+    """
+    if is_identity_matrix(matrix):
+        return image
+
+    height = int(np.round(output_shape[0]))
+    width = int(np.round(output_shape[1]))
+
+    cv2_matrix = matrix[:2, :]
+
+    warp_fn = maybe_process_in_chunks(
+        warp_affine_with_value_extension,
+        matrix=cv2_matrix,
+        dsize=(width, height),
+        flags=interpolation,
+        border_mode=border_mode,
+        border_value=fill,
+    )
+    return warp_fn(image)
+
+
+@handle_empty_array("keypoints")
+@angle_2pi_range
+def keypoints_affine(
+    keypoints: np.ndarray,
+    matrix: np.ndarray,
+    image_shape: tuple[int, int],
+    scale: dict[str, float],
+    border_mode: int,
+) -> np.ndarray:
+    """Apply an affine transformation to keypoints.
+
+    This function transforms keypoints using the given affine transformation matrix.
+    It handles reflection padding if necessary, updates coordinates, angles, and scales.
+
+    Args:
+        keypoints (np.ndarray): Array of keypoints with shape (N, 4+) where N is the number of keypoints.
+                                Each keypoint is represented as [x, y, angle, scale, ...].
+        matrix (np.ndarray): The 2x3 or 3x3 affine transformation matrix.
+        image_shape (tuple[int, int]): Shape of the image (height, width).
+        scale (dict[str, float]): Dictionary containing scale factors for x and y directions.
+                                  Expected keys are 'x' and 'y'.
+        border_mode (int): Border mode for handling keypoints near image edges.
+                            Use cv2.BORDER_REFLECT_101, cv2.BORDER_REFLECT, etc.
+
+    Returns:
+        np.ndarray: Transformed keypoints array with the same shape as input.
+
+    Notes:
+        - The function applies reflection padding if the mode is in REFLECT_BORDER_MODES.
+        - Coordinates (x, y) are transformed using the affine matrix.
+        - Angles are adjusted based on the rotation component of the affine transformation.
+        - Scales are multiplied by the maximum of x and y scale factors.
+        - The @angle_2pi_range decorator ensures angles remain in the [0, 2π] range.
+
+    Examples:
+        >>> keypoints = np.array([[100, 100, 0, 1]])
+        >>> matrix = np.array([[1.5, 0, 10], [0, 1.2, 20]])
+        >>> scale = {'x': 1.5, 'y': 1.2}
+        >>> transformed_keypoints = keypoints_affine(keypoints, matrix, (480, 640), scale, cv2.BORDER_REFLECT_101)
+
+    """
+    keypoints = keypoints.copy().astype(np.float32)
+
+    if is_identity_matrix(matrix):
+        return keypoints
+
+    if border_mode in REFLECT_BORDER_MODES:
+        # Step 1: Compute affine transform padding
+        pad_left, pad_right, pad_top, pad_bottom = calculate_affine_transform_padding(
+            matrix,
+            image_shape,
+        )
+        grid_dimensions = get_pad_grid_dimensions(
+            pad_top,
+            pad_bottom,
+            pad_left,
+            pad_right,
+            image_shape,
+        )
+        keypoints = generate_reflected_keypoints(
+            keypoints,
+            grid_dimensions,
+            image_shape,
+            center_in_origin=True,
+        )
+
+    # Extract x, y coordinates (z is preserved)
+    xy = keypoints[:, :2]
+
+    # Ensure matrix is 2x3
+    if matrix.shape == (3, 3):
+        matrix = matrix[:2]
+
+    # Transform x, y coordinates
+    xy_transformed = cv2.transform(xy.reshape(-1, 1, 2), matrix).squeeze()
+
+    # Calculate angle adjustment
+    angle_adjustment = rotation2d_matrix_to_euler_angles(matrix[:2, :2], y_up=False)
+
+    # Update angles (now at index 3)
+    keypoints[:, 3] = keypoints[:, 3] + angle_adjustment
+
+    # Update scales (now at index 4)
+    max_scale = max(scale["x"], scale["y"])
+    keypoints[:, 4] *= max_scale
+
+    # Update x, y coordinates and preserve z
+    keypoints[:, :2] = xy_transformed
+
+    return keypoints
+
+
+@handle_empty_array("points")
+def apply_affine_to_points(points: np.ndarray, matrix: np.ndarray) -> np.ndarray:
+    """Apply affine transformation to a set of points.
+
+    This function handles potential division by zero by replacing zero values
+    in the homogeneous coordinate with a small epsilon value.
+
+    Args:
+        points (np.ndarray): Array of points with shape (N, 2).
+        matrix (np.ndarray): 3x3 affine transformation matrix.
+
+    Returns:
+        np.ndarray: Transformed points with shape (N, 2).
+
+    """
+    homogeneous_points = np.column_stack([points, np.ones(points.shape[0])])
+    transformed_points = homogeneous_points @ matrix.T
+
+    # Handle potential division by zero
+    epsilon = np.finfo(transformed_points.dtype).eps
+    transformed_points[:, 2] = np.where(
+        np.abs(transformed_points[:, 2]) < epsilon,
+        np.sign(transformed_points[:, 2]) * epsilon,
+        transformed_points[:, 2],
+    )
+
+    return transformed_points[:, :2] / transformed_points[:, 2:]
+
+
+def calculate_affine_transform_padding(
+    matrix: np.ndarray,
+    image_shape: tuple[int, int],
+) -> tuple[int, int, int, int]:
+    """Calculate the necessary padding for an affine transformation to avoid empty spaces."""
+    height, width = image_shape[:2]
+
+    # Check for identity transform
+    if is_identity_matrix(matrix):
+        return (0, 0, 0, 0)
+
+    # Original corners
+    corners = np.array([[0, 0], [width, 0], [width, height], [0, height]])
+
+    # Transform corners
+    transformed_corners = apply_affine_to_points(corners, matrix)
+
+    # Ensure transformed_corners is 2D
+    transformed_corners = transformed_corners.reshape(-1, 2)
+
+    # Find box that includes both original and transformed corners
+    all_corners = np.vstack((corners, transformed_corners))
+    min_x, min_y = all_corners.min(axis=0)
+    max_x, max_y = all_corners.max(axis=0)
+
+    # Compute the inverse transform
+    inverse_matrix = np.linalg.inv(matrix)
+
+    # Apply inverse transform to all corners of the bounding box
+    bbox_corners = np.array(
+        [[min_x, min_y], [max_x, min_y], [max_x, max_y], [min_x, max_y]],
+    )
+    inverse_corners = apply_affine_to_points(bbox_corners, inverse_matrix).reshape(
+        -1,
+        2,
+    )
+
+    min_x, min_y = inverse_corners.min(axis=0)
+    max_x, max_y = inverse_corners.max(axis=0)
+
+    pad_left = max(0, math.ceil(0 - min_x))
+    pad_right = max(0, math.ceil(max_x - width))
+    pad_top = max(0, math.ceil(0 - min_y))
+    pad_bottom = max(0, math.ceil(max_y - height))
+
+    return pad_left, pad_right, pad_top, pad_bottom
+
+
+@handle_empty_array("bboxes")
+def bboxes_affine_largest_box(bboxes: np.ndarray, matrix: np.ndarray) -> np.ndarray:
+    """Apply an affine transformation to bounding boxes and return the largest enclosing boxes.
+
+    This function transforms each corner of every bounding box using the given affine transformation
+    matrix, then computes the new bounding boxes that fully enclose the transformed corners.
+
+    Args:
+        bboxes (np.ndarray): An array of bounding boxes with shape (N, 4+) where N is the number of
+                             bounding boxes. Each row should contain [x_min, y_min, x_max, y_max]
+                             followed by any additional attributes (e.g., class labels).
+        matrix (np.ndarray): The 3x3 affine transformation matrix to apply.
+
+    Returns:
+        np.ndarray: An array of transformed bounding boxes with the same shape as the input.
+                    Each row contains [new_x_min, new_y_min, new_x_max, new_y_max] followed by
+                    any additional attributes from the input bounding boxes.
+
+    Note:
+        - This function assumes that the input bounding boxes are in the format [x_min, y_min, x_max, y_max].
+        - The resulting bounding boxes are the smallest axis-aligned boxes that completely
+          enclose the transformed original boxes. They may be larger than the minimal possible
+          bounding box if the original box becomes rotated.
+        - Any additional attributes beyond the first 4 coordinates are preserved unchanged.
+        - This method is called "largest box" because it returns the largest axis-aligned box
+          that encloses all corners of the transformed bounding box.
+
+    Examples:
+        >>> bboxes = np.array([[10, 10, 20, 20, 1], [30, 30, 40, 40, 2]])  # Two boxes with class labels
+        >>> matrix = np.array([[2, 0, 5], [0, 2, 5], [0, 0, 1]])  # Scale by 2 and translate by (5, 5)
+        >>> transformed_bboxes = bboxes_affine_largest_box(bboxes, matrix)
+        >>> print(transformed_bboxes)
+        [[ 25.  25.  45.  45.   1.]
+         [ 65.  65.  85.  85.   2.]]
+
+    """
+    # Extract corners of all bboxes
+    x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]
+
+    corners = (
+        np.array([[x_min, y_min], [x_max, y_min], [x_max, y_max], [x_min, y_max]]).transpose(2, 0, 1).reshape(-1, 2)
+    )
+
+    # Transform all corners at once
+    transformed_corners = apply_affine_to_points(corners, matrix).reshape(-1, 4, 2)
+
+    # Compute new bounding boxes
+    new_x_min = np.min(transformed_corners[:, :, 0], axis=1)
+    new_x_max = np.max(transformed_corners[:, :, 0], axis=1)
+    new_y_min = np.min(transformed_corners[:, :, 1], axis=1)
+    new_y_max = np.max(transformed_corners[:, :, 1], axis=1)
+
+    return np.column_stack([new_x_min, new_y_min, new_x_max, new_y_max, bboxes[:, 4:]])
+
+
+@handle_empty_array("bboxes")
+def bboxes_affine_ellipse(bboxes: np.ndarray, matrix: np.ndarray) -> np.ndarray:
+    """Apply an affine transformation to bounding boxes using an ellipse approximation method.
+
+    This function transforms bounding boxes by approximating each box with an ellipse,
+    transforming points along the ellipse's circumference, and then computing the
+    new bounding box that encloses the transformed ellipse.
+
+    Args:
+        bboxes (np.ndarray): An array of bounding boxes with shape (N, 4+) where N is the number of
+                             bounding boxes. Each row should contain [x_min, y_min, x_max, y_max]
+                             followed by any additional attributes (e.g., class labels).
+        matrix (np.ndarray): The 3x3 affine transformation matrix to apply.
+
+    Returns:
+        np.ndarray: An array of transformed bounding boxes with the same shape as the input.
+                    Each row contains [new_x_min, new_y_min, new_x_max, new_y_max] followed by
+                    any additional attributes from the input bounding boxes.
+
+    Note:
+        - This function assumes that the input bounding boxes are in the format [x_min, y_min, x_max, y_max].
+        - The ellipse approximation method can provide a tighter bounding box compared to the
+          largest box method, especially for rotations.
+        - 360 points are used to approximate each ellipse, which provides a good balance between
+          accuracy and computational efficiency.
+        - Any additional attributes beyond the first 4 coordinates are preserved unchanged.
+        - This method may be more suitable for objects that are roughly elliptical in shape.
+
+    """
+    x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]
+    bbox_width = (x_max - x_min) / 2
+    bbox_height = (y_max - y_min) / 2
+    center_x = x_min + bbox_width
+    center_y = y_min + bbox_height
+
+    angles = np.arange(0, 360, dtype=np.float32)
+    cos_angles = np.cos(np.radians(angles))
+    sin_angles = np.sin(np.radians(angles))
+
+    # Generate points for all ellipses at once
+    x = bbox_width[:, np.newaxis] * sin_angles + center_x[:, np.newaxis]
+    y = bbox_height[:, np.newaxis] * cos_angles + center_y[:, np.newaxis]
+    points = np.stack([x, y], axis=-1).reshape(-1, 2)
+
+    # Transform all points at once using the helper function
+    transformed_points = apply_affine_to_points(points, matrix)
+
+    transformed_points = transformed_points.reshape(len(bboxes), -1, 2)
+
+    # Compute new bounding boxes
+    new_x_min = np.min(transformed_points[:, :, 0], axis=1)
+    new_x_max = np.max(transformed_points[:, :, 0], axis=1)
+    new_y_min = np.min(transformed_points[:, :, 1], axis=1)
+    new_y_max = np.max(transformed_points[:, :, 1], axis=1)
+
+    return np.column_stack([new_x_min, new_y_min, new_x_max, new_y_max, bboxes[:, 4:]])
+
+
+@handle_empty_array("bboxes")
+def bboxes_affine(
+    bboxes: np.ndarray,
+    matrix: np.ndarray,
+    rotate_method: Literal["largest_box", "ellipse"],
+    image_shape: tuple[int, int],
+    border_mode: int,
+    output_shape: tuple[int, int],
+) -> np.ndarray:
+    """Apply an affine transformation to bounding boxes.
+
+    For reflection border modes (cv2.BORDER_REFLECT_101, cv2.BORDER_REFLECT), this function:
+    1. Calculates necessary padding to avoid information loss
+    2. Applies padding to the bounding boxes
+    3. Adjusts the transformation matrix to account for padding
+    4. Applies the affine transformation
+    5. Validates the transformed bounding boxes
+
+    For other border modes, it directly applies the affine transformation without padding.
+
+    Args:
+        bboxes (np.ndarray): Input bounding boxes
+        matrix (np.ndarray): Affine transformation matrix
+        rotate_method (str): Method for rotating bounding boxes ('largest_box' or 'ellipse')
+        image_shape (Sequence[int]): Shape of the input image
+        border_mode (int): OpenCV border mode
+        output_shape (Sequence[int]): Shape of the output image
+
+    Returns:
+        np.ndarray: Transformed and normalized bounding boxes
+
+    """
+    if is_identity_matrix(matrix):
+        return bboxes
+
+    bboxes = denormalize_bboxes(bboxes, image_shape)
+
+    if border_mode in REFLECT_BORDER_MODES:
+        # Step 1: Compute affine transform padding
+        pad_left, pad_right, pad_top, pad_bottom = calculate_affine_transform_padding(
+            matrix,
+            image_shape,
+        )
+        grid_dimensions = get_pad_grid_dimensions(
+            pad_top,
+            pad_bottom,
+            pad_left,
+            pad_right,
+            image_shape,
+        )
+        bboxes = generate_reflected_bboxes(
+            bboxes,
+            grid_dimensions,
+            image_shape,
+            center_in_origin=True,
+        )
+
+    # Apply affine transform
+    if rotate_method == "largest_box":
+        transformed_bboxes = bboxes_affine_largest_box(bboxes, matrix)
+    elif rotate_method == "ellipse":
+        transformed_bboxes = bboxes_affine_ellipse(bboxes, matrix)
+    else:
+        raise ValueError(f"Method {rotate_method} is not a valid rotation method.")
+
+    # Validate and normalize bboxes
+    validated_bboxes = validate_bboxes(transformed_bboxes, output_shape)
+
+    return normalize_bboxes(validated_bboxes, output_shape)
+
+
+def to_distance_maps(
+    keypoints: np.ndarray,
+    image_shape: tuple[int, int],
+    inverted: bool = False,
+) -> np.ndarray:
+    """Generate a ``(H,W,N)`` array of distance maps for ``N`` keypoints.
+    The ``n``-th distance map contains at every location ``(y, x)`` the
+    euclidean distance to the ``n``-th keypoint.
+    This function can be used as a helper when augmenting keypoints with a
+    method that only supports the augmentation of images.
+
+    Args:
+        keypoints (np.ndarray): A numpy array of shape (N, 2+) where N is the number of keypoints.
+                   Each row represents a keypoint's (x, y) coordinates.
+        image_shape (tuple[int, int]): Shape of the image (height, width)
+        inverted (bool): If ``True``, inverted distance maps are returned where each
+            distance value d is replaced by ``d/(d+1)``, i.e. the distance
+            maps have values in the range ``(0.0, 1.0]`` with ``1.0`` denoting
+            exactly the position of the respective keypoint.
+
+    Returns:
+        np.ndarray: A float32 array of shape (H, W, N) containing ``N`` distance maps for ``N``
+            keypoints. Each location ``(y, x, n)`` in the array denotes the
+            euclidean distance at ``(y, x)`` to the ``n``-th keypoint.
+            If `inverted` is ``True``, the distance ``d`` is replaced
+            by ``d/(d+1)``. The height and width of the array match the
+            height and width in ``image_shape``.
+
+    """
+    height, width = image_shape[:2]
+    if len(keypoints) == 0:
+        return np.zeros((height, width, 0), dtype=np.float32)
+
+    # Create coordinate grids
+    yy, xx = np.mgrid[:height, :width]
+
+    # Convert keypoints to numpy array
+    keypoints_array = np.array(keypoints)
+
+    # Compute distances for all keypoints at once
+    distances = np.sqrt(
+        (xx[..., np.newaxis] - keypoints_array[:, 0]) ** 2 + (yy[..., np.newaxis] - keypoints_array[:, 1]) ** 2,
+    )
+
+    if inverted:
+        return (1 / (distances + 1)).astype(np.float32)
+    return distances.astype(np.float32)
+
+
+def validate_if_not_found_coords(
+    if_not_found_coords: Sequence[int] | dict[str, Any] | None,
+) -> tuple[bool, float, float]:
+    """Validate and process `if_not_found_coords` parameter."""
+    if if_not_found_coords is None:
+        return True, -1, -1
+    if isinstance(if_not_found_coords, (tuple, list)):
+        if len(if_not_found_coords) != PAIR:
+            msg = "Expected tuple/list 'if_not_found_coords' to contain exactly two entries."
+            raise ValueError(msg)
+        return False, if_not_found_coords[0], if_not_found_coords[1]
+    if isinstance(if_not_found_coords, dict):
+        return False, if_not_found_coords["x"], if_not_found_coords["y"]
+
+    msg = "Expected if_not_found_coords to be None, tuple, list, or dict."
+    raise ValueError(msg)
+
+
+def find_keypoint(
+    position: tuple[int, int],
+    distance_map: np.ndarray,
+    threshold: float | None,
+    inverted: bool,
+) -> tuple[float, float] | None:
+    """Determine if a valid keypoint can be found at the given position."""
+    y, x = position
+    value = distance_map[y, x]
+    if not inverted and threshold is not None and value >= threshold:
+        return None
+    if inverted and threshold is not None and value <= threshold:
+        return None
+    return float(x), float(y)
+
+
+def from_distance_maps(
+    distance_maps: np.ndarray,
+    inverted: bool,
+    if_not_found_coords: Sequence[int] | dict[str, Any] | None = None,
+    threshold: float | None = None,
+) -> np.ndarray:
+    """Convert distance maps back to keypoints coordinates.
+
+    This function is the inverse of `to_distance_maps`. It takes distance maps generated for a set of keypoints
+    and reconstructs the original keypoint coordinates. The function supports both regular and inverted distance maps,
+    and can handle cases where keypoints are not found or fall outside a specified threshold.
+
+    Args:
+        distance_maps (np.ndarray): A 3D numpy array of shape (height, width, nb_keypoints) containing
+            distance maps for each keypoint. Each channel represents the distance map for one keypoint.
+        inverted (bool): If True, treats the distance maps as inverted (where higher values indicate
+            closer proximity to keypoints). If False, treats them as regular distance maps (where lower
+            values indicate closer proximity).
+        if_not_found_coords (Sequence[int] | dict[str, Any] | None, optional): Coordinates to use for
+            keypoints that are not found or fall outside the threshold. Can be:
+            - None: Drop keypoints that are not found.
+            - Sequence of two integers: Use these as (x, y) coordinates for not found keypoints.
+            - Dict with 'x' and 'y' keys: Use these values for not found keypoints.
+            Defaults to None.
+        threshold (float | None, optional): A threshold value to determine valid keypoints. For inverted
+            maps, values >= threshold are considered valid. For regular maps, values <= threshold are
+            considered valid. If None, all keypoints are considered valid. Defaults to None.
+
+    Returns:
+        np.ndarray: A 2D numpy array of shape (nb_keypoints, 2) containing the (x, y) coordinates
+        of the reconstructed keypoints. If `drop_if_not_found` is True (derived from if_not_found_coords),
+        the output may have fewer rows than input keypoints.
+
+    Raises:
+        ValueError: If the input `distance_maps` is not a 3D array.
+
+    Notes:
+        - The function uses vectorized operations for improved performance, especially with large numbers of keypoints.
+        - When `threshold` is None, all keypoints are considered valid, and `if_not_found_coords` is not used.
+        - The function assumes that the input distance maps are properly normalized and scaled according to the
+          original image dimensions.
+
+    Examples:
+        >>> distance_maps = np.random.rand(100, 100, 3)  # 3 keypoints
+        >>> inverted = True
+        >>> if_not_found_coords = [0, 0]
+        >>> threshold = 0.5
+        >>> keypoints = from_distance_maps(distance_maps, inverted, if_not_found_coords, threshold)
+        >>> print(keypoints.shape)
+        (3, 2)
+
+    """
+    if distance_maps.ndim != NUM_MULTI_CHANNEL_DIMENSIONS:
+        msg = f"Expected three-dimensional input, got {distance_maps.ndim} dimensions and shape {distance_maps.shape}."
+        raise ValueError(msg)
+    height, width, nb_keypoints = distance_maps.shape
+
+    drop_if_not_found, if_not_found_x, if_not_found_y = validate_if_not_found_coords(
+        if_not_found_coords,
+    )
+
+    # Find the indices of max/min values for all keypoints at once
+    if inverted:
+        hitidx_flat = np.argmax(
+            distance_maps.reshape(height * width, nb_keypoints),
+            axis=0,
+        )
+    else:
+        hitidx_flat = np.argmin(
+            distance_maps.reshape(height * width, nb_keypoints),
+            axis=0,
+        )
+
+    # Convert flat indices to 2D coordinates
+    hitidx_y, hitidx_x = np.unravel_index(hitidx_flat, (height, width))
+
+    # Create keypoints array
+    keypoints = np.column_stack((hitidx_x, hitidx_y)).astype(float)
+
+    if threshold is not None:
+        # Check threshold condition
+        if inverted:
+            valid_mask = distance_maps[hitidx_y, hitidx_x, np.arange(nb_keypoints)] >= threshold
+        else:
+            valid_mask = distance_maps[hitidx_y, hitidx_x, np.arange(nb_keypoints)] <= threshold
+
+        if not drop_if_not_found:
+            # Replace invalid keypoints with if_not_found_coords
+            keypoints[~valid_mask] = [if_not_found_x, if_not_found_y]
+        else:
+            # Keep only valid keypoints
+            return keypoints[valid_mask]
+
+    return keypoints
+
+
+# Group elements for D4 symmetry group
+D4_GROUP_ELEMENTS = ["e", "r90", "r180", "r270", "v", "hvt", "h", "t"]
+
+
+def d4(img: np.ndarray, group_member: Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"]) -> np.ndarray:
+    """Applies a `D_4` symmetry group transformation to an image array.
+
+    This function manipulates an image using transformations such as rotations and flips,
+    corresponding to the `D_4` dihedral group symmetry operations.
+    Each transformation is identified by a unique group member code.
+
+    Args:
+        img (np.ndarray): The input image array to transform.
+        group_member (Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"]): A string identifier indicating
+            the specific transformation to apply. Valid codes include:
+            - 'e': Identity (no transformation).
+            - 'r90': Rotate 90 degrees counterclockwise.
+            - 'r180': Rotate 180 degrees.
+            - 'r270': Rotate 270 degrees counterclockwise.
+            - 'v': Vertical flip.
+            - 'hvt': Transpose over second diagonal
+            - 'h': Horizontal flip.
+            - 't': Transpose (reflect over the main diagonal).
+
+    Returns:
+        np.ndarray: The transformed image array.
+
+    """
+    # Execute the appropriate transformation
+    return D4_TRANSFORMATIONS[group_member](img)
+
+
+def transpose(img: np.ndarray) -> np.ndarray:
+    """Transposes the first two dimensions of an array of any dimensionality.
+    Retains the order of any additional dimensions.
+
+    Args:
+        img (np.ndarray): Input array.
+
+    Returns:
+        np.ndarray: Transposed array.
+
+    """
+    # Generate the new axes order
+    new_axes = list(range(img.ndim))
+    new_axes[0], new_axes[1] = 1, 0  # Swap the first two dimensions
+
+    # Transpose the array using the new axes order
+    return img.transpose(new_axes)
+
+
+D4_TRANSFORMATIONS = {
+    "e": lambda x: x,  # Identity transformation
+    "r90": lambda x: rot90(x, 1),  # Rotate 90 degrees
+    "r180": lambda x: rot90(x, 2),  # Rotate 180 degrees
+    "r270": lambda x: rot90(x, 3),  # Rotate 270 degrees
+    "v": vflip,  # Vertical flip
+    "hvt": lambda x: transpose(rot90(x, 2)),  # Reflect over anti-diagonal
+    "h": hflip,  # Horizontal flip
+    "t": transpose,  # Transpose (reflect over main diagonal)
+}
+
+
+def transpose_images(images: np.ndarray) -> np.ndarray:
+    """Transpose a batch of images.
+
+    Args:
+        images (np.ndarray): Batch of images to transpose with shape:
+            - (N, H, W) for grayscale images
+            - (N, H, W, C) for multi-channel images
+            where N is the batch size, H is height, W is width, C is channels
+
+    Returns:
+        np.ndarray: Transposed batch of images with shape:
+            - (N, W, H) for grayscale images
+            - (N, W, H, C) for multi-channel images
+
+    """
+    # Generate the new axes order
+    new_axes = list(range(images.ndim))
+    # Swap dimensions 1 and 2 (Height and Width), preserving batch dimension and channels
+    new_axes[1], new_axes[2] = 2, 1
+
+    # Transpose the array using the new axes order
+    return images.transpose(new_axes)
+
+
+def transpose_volumes(volumes: np.ndarray) -> np.ndarray:
+    """Transpose a batch of volumes.
+
+    Args:
+        volumes (np.ndarray): Batch of volumes to transpose with shape:
+            - (N, D, H, W) for grayscale volumes
+            - (N, D, H, W, C) for multi-channel volumes
+            where N is the batch size, D is depth, H is height, W is width, C is channels
+
+    Returns:
+        np.ndarray: Transposed batch of volumes with shape:
+            - (N, D, W, H) for grayscale volumes
+            - (N, D, W, H, C) for multi-channel volumes
+
+    """
+    # Generate the new axes order
+    new_axes = list(range(volumes.ndim))
+    # Swap dimensions 2 and 3 (Height and Width), preserving batch, depth and channels
+    new_axes[2], new_axes[3] = 3, 2
+
+    # Transpose the array using the new axes order
+    return volumes.transpose(new_axes)
+
+
+def rot90(img: np.ndarray, factor: Literal[0, 1, 2, 3]) -> np.ndarray:
+    """Rotate an image 90 degrees counterclockwise.
+
+    Args:
+        img (np.ndarray): The input image to rotate.
+        factor (Literal[0, 1, 2, 3]): The number of 90-degree rotations to apply.
+
+    Returns:
+        np.ndarray: The rotated image.
+
+    """
+    return np.rot90(img, factor)
+
+
+def rot90_images(images: np.ndarray, factor: Literal[0, 1, 2, 3]) -> np.ndarray:
+    """Rotate a batch of images 90 degrees counter-clockwise multiple times.
+
+    Args:
+        images (np.ndarray): Batch of images to rotate with shape:
+            - (N, H, W) for grayscale images
+            - (N, H, W, C) for multi-channel images
+            where N is the batch size, H is height, W is width, C is channels
+        factor (Literal[0, 1, 2, 3]): The number of 90-degree rotations to apply.
+
+    Returns:
+        np.ndarray: Rotated batch of images with shape:
+            - (N, W, H) for grayscale images when factor is 1 or 3
+            - (N, H, W) for grayscale images when factor is 0 or 2
+            - (N, W, H, C) for multi-channel images when factor is 1 or 3
+            - (N, H, W, C) for multi-channel images when factor is 0 or 2
+
+    """
+    # Axes 1 (height) and 2 (width) for rotation, preserving batch dimension
+    return np.rot90(images, k=factor, axes=(1, 2))
+
+
+@handle_empty_array("bboxes")
+def bboxes_vflip(bboxes: np.ndarray) -> np.ndarray:
+    """Flip bounding boxes vertically.
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (num_boxes, 4+)
+
+    Returns:
+        np.ndarray: Vertically flipped bounding boxes
+
+    """
+    flipped_bboxes = bboxes.copy()
+    flipped_bboxes[:, 1] = 1 - bboxes[:, 3]  # new y_min = 1 - y_max
+    flipped_bboxes[:, 3] = 1 - bboxes[:, 1]  # new y_max = 1 - y_min
+
+    return flipped_bboxes
+
+
+@handle_empty_array("bboxes")
+def bboxes_hflip(bboxes: np.ndarray) -> np.ndarray:
+    """Flip bounding boxes horizontally.
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (num_boxes, 4+)
+
+    Returns:
+        np.ndarray: Horizontally flipped bounding boxes
+
+    """
+    flipped_bboxes = bboxes.copy()
+    flipped_bboxes[:, 0] = 1 - bboxes[:, 2]  # new x_min = 1 - x_max
+    flipped_bboxes[:, 2] = 1 - bboxes[:, 0]  # new x_max = 1 - x_min
+
+    return flipped_bboxes
+
+
+@handle_empty_array("bboxes")
+def bboxes_transpose(bboxes: np.ndarray) -> np.ndarray:
+    """Transpose bounding boxes along the main diagonal.
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (num_boxes, 4+)
+
+    Returns:
+        np.ndarray: Transposed bounding boxes
+
+    """
+    transposed_bboxes = bboxes.copy()
+    transposed_bboxes[:, [0, 1, 2, 3]] = bboxes[:, [1, 0, 3, 2]]
+
+    return transposed_bboxes
+
+
+@handle_empty_array("keypoints")
+@angle_2pi_range
+def keypoints_vflip(keypoints: np.ndarray, rows: int) -> np.ndarray:
+    """Flip keypoints vertically.
+
+    Args:
+        keypoints (np.ndarray): Array of keypoints with shape (num_keypoints, 2+)
+        rows (int): Number of rows in the image
+
+    Returns:
+        np.ndarray: Vertically flipped keypoints
+
+    """
+    flipped_keypoints = keypoints.copy().astype(np.float32)
+
+    # Flip y-coordinates
+    flipped_keypoints[:, 1] = (rows - 1) - keypoints[:, 1]
+
+    # Negate angles
+    flipped_keypoints[:, 3] = -keypoints[:, 3]
+
+    return flipped_keypoints
+
+
+@handle_empty_array("keypoints")
+@angle_2pi_range
+def keypoints_hflip(keypoints: np.ndarray, cols: int) -> np.ndarray:
+    """Flip keypoints horizontally.
+
+    Args:
+        keypoints (np.ndarray): Array of keypoints with shape (num_keypoints, 2+)
+        cols (int): Number of columns in the image
+
+    Returns:
+        np.ndarray: Horizontally flipped keypoints
+
+    """
+    flipped_keypoints = keypoints.copy().astype(np.float32)
+
+    # Flip x-coordinates
+    flipped_keypoints[:, 0] = (cols - 1) - keypoints[:, 0]
+
+    # Adjust angles
+    flipped_keypoints[:, 3] = np.pi - keypoints[:, 3]
+
+    return flipped_keypoints
+
+
+@handle_empty_array("keypoints")
+@angle_2pi_range
+def keypoints_transpose(keypoints: np.ndarray) -> np.ndarray:
+    """Transpose keypoints along the main diagonal.
+
+    Args:
+        keypoints (np.ndarray): Array of keypoints with shape (num_keypoints, 2+)
+
+    Returns:
+        np.ndarray: Transposed keypoints
+
+    """
+    transposed_keypoints = keypoints.copy()
+
+    # Swap x and y coordinates
+    transposed_keypoints[:, [0, 1]] = keypoints[:, [1, 0]]
+
+    # Adjust angles to reflect the coordinate swap
+    angles = keypoints[:, 3]
+    transposed_keypoints[:, 3] = np.where(
+        angles <= np.pi,
+        np.pi / 2 - angles,
+        3 * np.pi / 2 - angles,
+    )
+
+    return transposed_keypoints
+
+
+@preserve_channel_dim
+def pad(
+    img: np.ndarray,
+    min_height: int,
+    min_width: int,
+    border_mode: int,
+    value: tuple[float, ...] | float | None,
+) -> np.ndarray:
+    """Pad an image to ensure minimum dimensions.
+
+    This function adds padding to an image if its dimensions are smaller than
+    the specified minimum dimensions. Padding is added evenly on all sides.
+
+    Args:
+        img (np.ndarray): Input image to pad.
+        min_height (int): Minimum height of the output image.
+        min_width (int): Minimum width of the output image.
+        border_mode (int): OpenCV border mode for padding.
+        value (tuple[float, ...] | float | None): Value(s) to fill the border pixels.
+
+    Returns:
+        np.ndarray: Padded image with dimensions at least (min_height, min_width).
+
+    """
+    height, width = img.shape[:2]
+
+    if height < min_height:
+        h_pad_top = int((min_height - height) / 2.0)
+        h_pad_bottom = min_height - height - h_pad_top
+    else:
+        h_pad_top = 0
+        h_pad_bottom = 0
+
+    if width < min_width:
+        w_pad_left = int((min_width - width) / 2.0)
+        w_pad_right = min_width - width - w_pad_left
+    else:
+        w_pad_left = 0
+        w_pad_right = 0
+
+    img = pad_with_params(
+        img,
+        h_pad_top,
+        h_pad_bottom,
+        w_pad_left,
+        w_pad_right,
+        border_mode,
+        value,
+    )
+
+    if img.shape[:2] != (max(min_height, height), max(min_width, width)):
+        raise RuntimeError(
+            f"Invalid result shape. Got: {img.shape[:2]}. Expected: {(max(min_height, height), max(min_width, width))}",
+        )
+
+    return img
+
+
+def extend_value(value: tuple[float, ...] | float, num_channels: int) -> Sequence[float]:
+    """Extend value to a sequence of floats.
+
+    This function extends a value to a sequence of floats.
+    It is used to pad an image with a given value.
+
+    Args:
+        value (tuple[float, ...] | float): The value to extend.
+        num_channels (int): The number of channels in the image.
+
+    Returns:
+        Sequence[float]: The extended value.
+
+    """
+    return [value] * num_channels if isinstance(value, float) else value
+
+
+def copy_make_border_with_value_extension(
+    img: np.ndarray,
+    top: int,
+    bottom: int,
+    left: int,
+    right: int,
+    border_mode: int,
+    value: tuple[float, ...] | float,
+) -> np.ndarray:
+    """Copy and make border with value extension.
+
+    This function copies and makes border with value extension.
+    It is used to pad an image with a given value.
+
+    Args:
+        img (np.ndarray): The image to pad.
+        top (int): The amount to pad the top of the image.
+        bottom (int): The amount to pad the bottom of the image.
+        left (int): The amount to pad the left of the image.
+        right (int): The amount to pad the right of the image.
+        border_mode (int): The border mode to use.
+        value (tuple[float, ...] | float): The value to pad the image with.
+
+    Returns:
+        np.ndarray: The padded image.
+
+    """
+    # For 0-channel images, return empty array of correct padded size
+    if img.size == 0:
+        height, width = img.shape[:2]
+        return np.zeros(
+            (height + top + bottom, width + left + right, 0),
+            dtype=img.dtype,
+        )
+
+    num_channels = get_num_channels(img)
+    extended_value = extend_value(value, num_channels)
+
+    return cv2.copyMakeBorder(
+        img,
+        top,
+        bottom,
+        left,
+        right,
+        borderType=border_mode,
+        value=extended_value,
+    )
+
+
+@preserve_channel_dim
+def pad_with_params(
+    img: np.ndarray,
+    h_pad_top: int,
+    h_pad_bottom: int,
+    w_pad_left: int,
+    w_pad_right: int,
+    border_mode: int,
+    value: tuple[float, ...] | float | None,
+) -> np.ndarray:
+    """Pad an image with explicitly defined padding on each side.
+
+    This function adds specified amounts of padding to each side of the image.
+
+    Args:
+        img (np.ndarray): Input image to pad.
+        h_pad_top (int): Number of pixels to add at the top.
+        h_pad_bottom (int): Number of pixels to add at the bottom.
+        w_pad_left (int): Number of pixels to add on the left.
+        w_pad_right (int): Number of pixels to add on the right.
+        border_mode (int): OpenCV border mode for padding.
+        value (tuple[float, ...] | float | None): Value(s) to fill the border pixels.
+
+    Returns:
+        np.ndarray: Padded image.
+
+    """
+    pad_fn = maybe_process_in_chunks(
+        copy_make_border_with_value_extension,
+        top=h_pad_top,
+        bottom=h_pad_bottom,
+        left=w_pad_left,
+        right=w_pad_right,
+        border_mode=border_mode,
+        value=value,
+    )
+
+    return pad_fn(img)
+
+
+def pad_images_with_params(
+    images: np.ndarray,
+    h_pad_top: int,
+    h_pad_bottom: int,
+    w_pad_left: int,
+    w_pad_right: int,
+    border_mode: int,
+    value: tuple[float, ...] | float | None,
+) -> np.ndarray:
+    """Pad a batch of images with explicitly defined padding on each side.
+
+    This function adds specified amounts of padding to each side of the image for each
+    image in the batch.
+
+    Args:
+        images (np.ndarray): Input batch of images to pad.
+        h_pad_top (int): Number of pixels to add at the top.
+        h_pad_bottom (int): Number of pixels to add at the bottom.
+        w_pad_left (int): Number of pixels to add on the left.
+        w_pad_right (int): Number of pixels to add on the right.
+        border_mode (int): OpenCV border mode for padding.
+        value (tuple[float, ...] | float | None): Value(s) to fill the border pixels.
+
+    Returns:
+        np.ndarray: Padded batch of images.
+
+    """
+    no_channel_dim = images.ndim == 3
+    if no_channel_dim:
+        images = images[..., np.newaxis]
+
+    cv2np_border_modes = {
+        cv2.BORDER_CONSTANT: "constant",
+        cv2.BORDER_REPLICATE: "edge",
+        cv2.BORDER_REFLECT: "symmetric",
+        cv2.BORDER_WRAP: "wrap",
+        cv2.BORDER_REFLECT_101: "reflect",
+        cv2.BORDER_REFLECT101: "reflect",
+        cv2.BORDER_DEFAULT: "reflect",  # same as cv2.BORDER_REFLECT_101
+    }
+    mode = cv2np_border_modes[border_mode]
+
+    pad_width = ((0, 0), (h_pad_top, h_pad_bottom), (w_pad_left, w_pad_right), (0, 0))
+    if mode == "constant":
+        constant_values = np.array(((0, 0), (value, value), (value, value), (0, 0)), dtype=object)
+        kwargs = {"constant_values": constant_values}
+    else:
+        kwargs = {}
+
+    images = np.pad(images, pad_width=pad_width, mode=mode, **kwargs)
+    if no_channel_dim:
+        images = images[..., 0]
+
+    return images
+
+
+@preserve_channel_dim
+def remap(
+    img: np.ndarray,
+    map_x: np.ndarray,
+    map_y: np.ndarray,
+    interpolation: int,
+    border_mode: int,
+    value: tuple[float, ...] | float | None = None,
+) -> np.ndarray:
+    """Remap an image according to given coordinate maps.
+
+    This function applies a generic geometrical transformation using
+    mapping functions that specify the position of each pixel in the output image.
+
+    Args:
+        img (np.ndarray): Input image to transform.
+        map_x (np.ndarray): Map of x-coordinates with same height and width as the input image.
+        map_y (np.ndarray): Map of y-coordinates with same height and width as the input image.
+        interpolation (int): Interpolation method for resampling.
+        border_mode (int): OpenCV border mode for handling pixels outside the image boundaries.
+        value (tuple[float, ...] | float | None, optional): Border value(s) if border_mode is BORDER_CONSTANT.
+
+    Returns:
+        np.ndarray: Remapped image with the same shape as the input image.
+
+    """
+    # Combine map_x and map_y into a single map array of type CV_32FC2
+    map_xy = np.stack([map_x, map_y], axis=-1).astype(np.float32)
+
+    # Create remap function with chunks processing
+    remap_func = maybe_process_in_chunks(
+        cv2.remap,
+        map1=map_xy,
+        map2=None,
+        interpolation=interpolation,
+        borderMode=border_mode,
+        borderValue=value,
+    )
+
+    # Apply the remapping
+    return remap_func(img)
+
+
+def remap_keypoints_via_mask(
+    keypoints: np.ndarray,
+    map_x: np.ndarray,
+    map_y: np.ndarray,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Remap keypoints using mask and cv2.remap method."""
+    height, width = image_shape[:2]
+
+    # Handle empty keypoints array
+    if len(keypoints) == 0:
+        return np.zeros((0, 2 if keypoints.size == 0 else keypoints.shape[1]))
+
+    # Create mask where each keypoint has unique index
+    kp_mask = np.zeros((height, width), dtype=np.int16)
+    for idx, kp in enumerate(keypoints, start=1):
+        x, y = round(kp[0]), round(kp[1])
+        if 0 <= x < width and 0 <= y < height:
+            # Note: cv2.circle takes (x,y) coordinates
+            cv2.circle(kp_mask, (x, y), 1, idx, -1)
+
+    # Remap the mask
+    transformed_kp_mask = cv2.remap(
+        kp_mask,
+        map_x.astype(np.float32),
+        map_y.astype(np.float32),
+        cv2.INTER_NEAREST,
+    )
+
+    # Extract transformed keypoints
+    new_points = []
+    for idx, kp in enumerate(keypoints, start=1):
+        # Find points with this index
+        points = np.where(transformed_kp_mask == idx)
+        if len(points[0]) > 0:
+            # Convert back to (x,y) coordinates
+            new_points.append(np.concatenate([[points[1][0], points[0][0]], kp[2:]]))
+
+    return np.array(new_points) if new_points else np.zeros((0, keypoints.shape[1]))
+
+
+@handle_empty_array("keypoints")
+def remap_keypoints(
+    keypoints: np.ndarray,
+    map_x: np.ndarray,
+    map_y: np.ndarray,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Transform keypoints using coordinate mapping functions.
+
+    This function applies the inverse of the mapping defined by map_x and map_y
+    to keypoint coordinates. The inverse mapping is necessary because the mapping
+    functions define how pixels move from the source to the destination image,
+    while keypoints need to be transformed from the destination back to the source.
+
+    Args:
+        keypoints (np.ndarray): Array of keypoints with shape (N, 2+), where
+            the first two columns are x and y coordinates.
+        map_x (np.ndarray): Map of x-coordinates with shape equal to image_shape.
+        map_y (np.ndarray): Map of y-coordinates with shape equal to image_shape.
+        image_shape (tuple[int, int]): Shape (height, width) of the original image.
+
+    Returns:
+        np.ndarray: Transformed keypoints with the same shape as the input keypoints.
+            Returns an empty array if input keypoints is empty.
+
+    """
+    height, width = image_shape[:2]
+
+    # Extract x and y coordinates
+    x, y = keypoints[:, 0], keypoints[:, 1]
+
+    # Clip coordinates to image boundaries
+    x = np.clip(x, 0, width - 1)
+    y = np.clip(y, 0, height - 1)
+
+    # Convert to integer indices
+    x_idx, y_idx = x.astype(int), y.astype(int)
+    inv_map_x, inv_map_y = generate_inverse_distortion_map(map_x, map_y, image_shape[:2])
+    # Apply the inverse mapping
+    new_x = inv_map_x[y_idx, x_idx]
+    new_y = inv_map_y[y_idx, x_idx]
+
+    # Clip the new coordinates to ensure they're within the image bounds
+    new_x = np.clip(new_x, 0, width - 1)
+    new_y = np.clip(new_y, 0, height - 1)
+
+    # Create the transformed keypoints array
+    return np.column_stack([new_x, new_y, keypoints[:, 2:]])
+
+
+def generate_inverse_distortion_map(
+    map_x: np.ndarray,
+    map_y: np.ndarray,
+    shape: tuple[int, int],
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate inverse mapping for strong distortions."""
+    h, w = shape
+
+    # Initialize inverse maps
+    inv_map_x = np.zeros((h, w), dtype=np.float32)
+    inv_map_y = np.zeros((h, w), dtype=np.float32)
+
+    # For each source point, record where it maps to
+    for y in range(h):
+        for x in range(w):
+            # Get destination point
+            dst_x = map_x[y, x]
+            dst_y = map_y[y, x]
+
+            # If destination is within bounds
+            if 0 <= dst_x < w and 0 <= dst_y < h:
+                # Get neighborhood coordinates
+                dst_x_floor = int(np.floor(dst_x))
+                dst_x_ceil = min(dst_x_floor + 1, w - 1)
+                dst_y_floor = int(np.floor(dst_y))
+                dst_y_ceil = min(dst_y_floor + 1, h - 1)
+
+                # Fill neighborhood
+                for ny in range(dst_y_floor, dst_y_ceil + 1):
+                    for nx in range(dst_x_floor, dst_x_ceil + 1):
+                        # Only update if empty or closer to pixel center
+                        if inv_map_x[ny, nx] == 0 or (
+                            abs(nx - dst_x) + abs(ny - dst_y)
+                            < abs(nx - inv_map_x[ny, nx]) + abs(ny - inv_map_y[ny, nx])
+                        ):
+                            inv_map_x[ny, nx] = x
+                            inv_map_y[ny, nx] = y
+
+    return inv_map_x, inv_map_y
+
+
+@handle_empty_array("bboxes")
+def remap_bboxes(
+    bboxes: np.ndarray,
+    map_x: np.ndarray,
+    map_y: np.ndarray,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Remap bounding boxes using displacement maps."""
+    # Convert bboxes to mask
+    bbox_masks = bboxes_to_mask(bboxes, image_shape)
+
+    # Ensure maps are float32
+    map_x = map_x.astype(np.float32)
+    map_y = map_y.astype(np.float32)
+
+    transformed_masks = remap(bbox_masks, map_x, map_y, cv2.INTER_NEAREST, cv2.BORDER_CONSTANT, value=0)
+
+    # Convert masks back to bboxes
+    return mask_to_bboxes(transformed_masks, bboxes)
+
+
+def generate_displacement_fields(
+    image_shape: tuple[int, int],
+    alpha: float,
+    sigma: float,
+    same_dxdy: bool,
+    kernel_size: tuple[int, int],
+    random_generator: np.random.Generator,
+    noise_distribution: Literal["gaussian", "uniform"],
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate displacement fields for elastic transform.
+
+    This function generates displacement fields for elastic transform based on the provided parameters.
+    It generates noise either from a Gaussian or uniform distribution and normalizes it to the range [-1, 1].
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+        alpha (float): The alpha parameter for the elastic transform.
+        sigma (float): The sigma parameter for the elastic transform.
+        same_dxdy (bool): Whether to use the same displacement field for both x and y directions.
+        kernel_size (tuple[int, int]): The size of the kernel for the elastic transform.
+        random_generator (np.random.Generator): The random number generator to use.
+        noise_distribution (Literal["gaussian", "uniform"]): The distribution of the noise.
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: A tuple containing:
+            - fields: The displacement fields for the elastic transform.
+            - output_shape: The output shape of the elastic warp.
+
+    """
+    # Pre-allocate memory and generate noise in one step
+    if noise_distribution == "gaussian":
+        # Generate and normalize in one step, directly as float32
+        fields = random_generator.standard_normal(
+            (1 if same_dxdy else 2, *image_shape[:2]),
+            dtype=np.float32,
+        )
+        # Normalize inplace
+        max_abs = np.abs(fields, out=np.empty_like(fields)).max()
+        if max_abs > 1e-6:
+            fields /= max_abs
+    else:  # uniform is already normalized to [-1, 1]
+        fields = random_generator.uniform(
+            -1,
+            1,
+            size=(1 if same_dxdy else 2, *image_shape[:2]),
+        ).astype(np.float32)
+
+    # # Apply Gaussian blur if needed using fast OpenCV operations
+    # When kernel_size is (0,0) cv2.GaussianBlur uses automatic kernel size. Kernel == (0,0) is NOT a noop.
+    # Reshape to 2D array (combining first dimension with height)
+    shape = fields.shape
+    fields = fields.reshape(-1, shape[-1])
+
+    # Apply blur to all fields at once
+    cv2.GaussianBlur(
+        fields,
+        kernel_size,
+        sigma,
+        dst=fields,
+        borderType=cv2.BORDER_REPLICATE,
+    )
+
+    # Restore original shape
+    fields = fields.reshape(shape)
+
+    # Scale by alpha inplace
+    fields *= alpha
+
+    # Return views of the array to avoid copies
+    return (fields[0], fields[0]) if same_dxdy else (fields[0], fields[1])
+
+
+@handle_empty_array("bboxes")
+def pad_bboxes(
+    bboxes: np.ndarray,
+    pad_top: int,
+    pad_bottom: int,
+    pad_left: int,
+    pad_right: int,
+    border_mode: int,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Pad bounding boxes by a given amount.
+
+    This function pads bounding boxes by a given amount.
+    It handles both reflection and padding.
+
+    Args:
+        bboxes (np.ndarray): The bounding boxes to pad.
+        pad_top (int): The amount to pad the top of the bounding boxes.
+        pad_bottom (int): The amount to pad the bottom of the bounding boxes.
+        pad_left (int): The amount to pad the left of the bounding boxes.
+        pad_right (int): The amount to pad the right of the bounding boxes.
+        border_mode (int): The border mode to use.
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: The padded bounding boxes.
+
+    """
+    if border_mode not in REFLECT_BORDER_MODES:
+        shift_vector = np.array([pad_left, pad_top, pad_left, pad_top])
+        return shift_bboxes(bboxes, shift_vector)
+
+    grid_dimensions = get_pad_grid_dimensions(
+        pad_top,
+        pad_bottom,
+        pad_left,
+        pad_right,
+        image_shape,
+    )
+
+    bboxes = generate_reflected_bboxes(bboxes, grid_dimensions, image_shape)
+
+    # Calculate the number of grid cells added on each side
+    original_row, original_col = grid_dimensions["original_position"]
+
+    image_height, image_width = image_shape[:2]
+
+    # Subtract the offset based on the number of added grid cells
+    left_shift = original_col * image_width - pad_left
+    top_shift = original_row * image_height - pad_top
+
+    shift_vector = np.array([-left_shift, -top_shift, -left_shift, -top_shift])
+
+    bboxes = shift_bboxes(bboxes, shift_vector)
+
+    new_height = pad_top + pad_bottom + image_height
+    new_width = pad_left + pad_right + image_width
+
+    return validate_bboxes(bboxes, (new_height, new_width))
+
+
+def validate_bboxes(bboxes: np.ndarray, image_shape: Sequence[int]) -> np.ndarray:
+    """Validate bounding boxes and remove invalid ones.
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (n, 4) where each row is [x_min, y_min, x_max, y_max].
+        image_shape (tuple[int, int]): Shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: Array of valid bounding boxes, potentially with fewer boxes than the input.
+
+    Examples:
+        >>> bboxes = np.array([[10, 20, 30, 40], [-10, -10, 5, 5], [100, 100, 120, 120]])
+        >>> valid_bboxes = validate_bboxes(bboxes, (100, 100))
+        >>> print(valid_bboxes)
+        [[10 20 30 40]]
+
+    """
+    rows, cols = image_shape[:2]
+
+    x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]
+
+    valid_indices = (x_max > 0) & (y_max > 0) & (x_min < cols) & (y_min < rows)
+
+    return bboxes[valid_indices]
+
+
+def shift_bboxes(bboxes: np.ndarray, shift_vector: np.ndarray) -> np.ndarray:
+    """Shift bounding boxes by a given vector.
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (n, m) where n is the number of bboxes
+                             and m >= 4. The first 4 columns are [x_min, y_min, x_max, y_max].
+        shift_vector (np.ndarray): Vector to shift the bounding boxes by, with shape (4,) for
+                                   [shift_x, shift_y, shift_x, shift_y].
+
+    Returns:
+        np.ndarray: Shifted bounding boxes with the same shape as input.
+
+    """
+    # Create a copy of the input array to avoid modifying it in-place
+    shifted_bboxes = bboxes.copy()
+
+    # Add the shift vector to the first 4 columns
+    shifted_bboxes[:, :4] += shift_vector
+
+    return shifted_bboxes
+
+
+def get_pad_grid_dimensions(
+    pad_top: int,
+    pad_bottom: int,
+    pad_left: int,
+    pad_right: int,
+    image_shape: tuple[int, int],
+) -> dict[str, tuple[int, int]]:
+    """Calculate the dimensions of the grid needed for reflection padding and the position of the original image.
+
+    Args:
+        pad_top (int): Number of pixels to pad above the image.
+        pad_bottom (int): Number of pixels to pad below the image.
+        pad_left (int): Number of pixels to pad to the left of the image.
+        pad_right (int): Number of pixels to pad to the right of the image.
+        image_shape (tuple[int, int]): Shape of the original image as (height, width).
+
+    Returns:
+        dict[str, tuple[int, int]]: A dictionary containing:
+            - 'grid_shape': A tuple (grid_rows, grid_cols) where:
+                - grid_rows (int): Number of times the image needs to be repeated vertically.
+                - grid_cols (int): Number of times the image needs to be repeated horizontally.
+            - 'original_position': A tuple (original_row, original_col) where:
+                - original_row (int): Row index of the original image in the grid.
+                - original_col (int): Column index of the original image in the grid.
+
+    """
+    rows, cols = image_shape[:2]
+
+    grid_rows = 1 + math.ceil(pad_top / rows) + math.ceil(pad_bottom / rows)
+    grid_cols = 1 + math.ceil(pad_left / cols) + math.ceil(pad_right / cols)
+    original_row = math.ceil(pad_top / rows)
+    original_col = math.ceil(pad_left / cols)
+
+    return {
+        "grid_shape": (grid_rows, grid_cols),
+        "original_position": (original_row, original_col),
+    }
+
+
+def generate_reflected_bboxes(
+    bboxes: np.ndarray,
+    grid_dims: dict[str, tuple[int, int]],
+    image_shape: tuple[int, int],
+    center_in_origin: bool = False,
+) -> np.ndarray:
+    """Generate reflected bounding boxes for the entire reflection grid.
+
+    Args:
+        bboxes (np.ndarray): Original bounding boxes.
+        grid_dims (dict[str, tuple[int, int]]): Grid dimensions and original position.
+        image_shape (tuple[int, int]): Shape of the original image as (height, width).
+        center_in_origin (bool): If True, center the grid at the origin. Default is False.
+
+    Returns:
+        np.ndarray: Array of reflected and shifted bounding boxes for the entire grid.
+
+    """
+    rows, cols = image_shape[:2]
+    grid_rows, grid_cols = grid_dims["grid_shape"]
+    original_row, original_col = grid_dims["original_position"]
+
+    # Prepare flipped versions of bboxes
+    bboxes_hflipped = flip_bboxes(bboxes, flip_horizontal=True, image_shape=image_shape)
+    bboxes_vflipped = flip_bboxes(bboxes, flip_vertical=True, image_shape=image_shape)
+    bboxes_hvflipped = flip_bboxes(
+        bboxes,
+        flip_horizontal=True,
+        flip_vertical=True,
+        image_shape=image_shape,
+    )
+
+    # Shift all versions to the original position
+    shift_vector = np.array(
+        [
+            original_col * cols,
+            original_row * rows,
+            original_col * cols,
+            original_row * rows,
+        ],
+    )
+    bboxes = shift_bboxes(bboxes, shift_vector)
+    bboxes_hflipped = shift_bboxes(bboxes_hflipped, shift_vector)
+    bboxes_vflipped = shift_bboxes(bboxes_vflipped, shift_vector)
+    bboxes_hvflipped = shift_bboxes(bboxes_hvflipped, shift_vector)
+
+    new_bboxes = []
+
+    for grid_row in range(grid_rows):
+        for grid_col in range(grid_cols):
+            # Determine which version of bboxes to use based on grid position
+            if (grid_row - original_row) % 2 == 0 and (grid_col - original_col) % 2 == 0:
+                current_bboxes = bboxes
+            elif (grid_row - original_row) % 2 == 0:
+                current_bboxes = bboxes_hflipped
+            elif (grid_col - original_col) % 2 == 0:
+                current_bboxes = bboxes_vflipped
+            else:
+                current_bboxes = bboxes_hvflipped
+
+            # Shift to the current grid cell
+            cell_shift = np.array(
+                [
+                    (grid_col - original_col) * cols,
+                    (grid_row - original_row) * rows,
+                    (grid_col - original_col) * cols,
+                    (grid_row - original_row) * rows,
+                ],
+            )
+            shifted_bboxes = shift_bboxes(current_bboxes, cell_shift)
+
+            new_bboxes.append(shifted_bboxes)
+
+    result = np.vstack(new_bboxes)
+
+    return shift_bboxes(result, -shift_vector) if center_in_origin else result
+
+
+@handle_empty_array("bboxes")
+def flip_bboxes(
+    bboxes: np.ndarray,
+    flip_horizontal: bool = False,
+    flip_vertical: bool = False,
+    image_shape: tuple[int, int] = (0, 0),
+) -> np.ndarray:
+    """Flip bounding boxes horizontally and/or vertically.
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (n, m) where each row is
+            [x_min, y_min, x_max, y_max, ...].
+        flip_horizontal (bool): Whether to flip horizontally.
+        flip_vertical (bool): Whether to flip vertically.
+        image_shape (tuple[int, int]): Shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: Flipped bounding boxes.
+
+    """
+    rows, cols = image_shape[:2]
+    flipped_bboxes = bboxes.copy()
+    if flip_horizontal:
+        flipped_bboxes[:, [0, 2]] = cols - flipped_bboxes[:, [2, 0]]
+    if flip_vertical:
+        flipped_bboxes[:, [1, 3]] = rows - flipped_bboxes[:, [3, 1]]
+    return flipped_bboxes
+
+
+@preserve_channel_dim
+def distort_image(
+    image: np.ndarray,
+    generated_mesh: np.ndarray,
+    interpolation: int,
+) -> np.ndarray:
+    """Apply perspective distortion to an image based on a generated mesh.
+
+    This function applies a perspective transformation to each cell of the image defined by the
+    generated mesh. The distortion is applied using OpenCV's perspective transformation and
+    blending techniques.
+
+    Args:
+        image (np.ndarray): The input image to be distorted. Can be a 2D grayscale image or a
+                            3D color image.
+        generated_mesh (np.ndarray): A 2D array where each row represents a quadrilateral cell
+                                    as [x1, y1, x2, y2, dst_x1, dst_y1, dst_x2, dst_y2, dst_x3, dst_y3, dst_x4, dst_y4].
+                                    The first four values define the source rectangle, and the last eight values
+                                    define the destination quadrilateral.
+        interpolation (int): Interpolation method to be used in the perspective transformation.
+                             Should be one of the OpenCV interpolation flags (e.g., cv2.INTER_LINEAR).
+
+    Returns:
+        np.ndarray: The distorted image with the same shape and dtype as the input image.
+
+    Note:
+        - The function preserves the channel dimension of the input image.
+        - Each cell of the generated mesh is transformed independently and then blended into the output image.
+        - The distortion is applied using perspective transformation, which allows for more complex
+          distortions compared to affine transformations.
+
+    Examples:
+        >>> image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+        >>> mesh = np.array([[0, 0, 50, 50, 5, 5, 45, 5, 45, 45, 5, 45]])
+        >>> distorted = distort_image(image, mesh, cv2.INTER_LINEAR)
+        >>> distorted.shape
+        (100, 100, 3)
+
+    """
+    distorted_image = np.zeros_like(image)
+
+    for mesh in generated_mesh:
+        # Extract source rectangle and destination quadrilateral
+        x1, y1, x2, y2 = mesh[:4]  # Source rectangle
+        dst_quad = mesh[4:].reshape(4, 2)  # Destination quadrilateral
+
+        # Convert source rectangle to quadrilateral
+        src_quad = np.array(
+            [
+                [x1, y1],  # Top-left
+                [x2, y1],  # Top-right
+                [x2, y2],  # Bottom-right
+                [x1, y2],  # Bottom-left
+            ],
+            dtype=np.float32,
+        )
+
+        # Calculate Perspective transformation matrix
+        perspective_mat = cv2.getPerspectiveTransform(src_quad, dst_quad)
+
+        # Apply Perspective transformation
+        warped = cv2.warpPerspective(
+            image,
+            perspective_mat,
+            (image.shape[1], image.shape[0]),
+            flags=interpolation,
+        )
+
+        # Create mask for the transformed region
+        mask = np.zeros(image.shape[:2], dtype=np.uint8)
+        cv2.fillConvexPoly(mask, np.int32(dst_quad), 255)
+
+        # Copy only the warped quadrilateral area to the output image
+        distorted_image = cv2.copyTo(warped, mask, distorted_image)
+
+    return distorted_image
+
+
+@handle_empty_array("bboxes")
+def bbox_distort_image(
+    bboxes: np.ndarray,
+    generated_mesh: np.ndarray,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Distort bounding boxes based on a generated mesh.
+
+    This function applies a perspective transformation to each bounding box based on the provided generated mesh.
+    It ensures that the bounding boxes are clipped to the image boundaries after transformation.
+
+    Args:
+        bboxes (np.ndarray): The bounding boxes to distort.
+        generated_mesh (np.ndarray): The generated mesh to distort the bounding boxes with.
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: The distorted bounding boxes.
+
+    """
+    bboxes = bboxes.copy()
+    masks = masks_from_bboxes(bboxes, image_shape)
+
+    transformed_masks = cv2.merge(
+        [distort_image(mask, generated_mesh, cv2.INTER_NEAREST) for mask in masks],
+    )
+
+    if transformed_masks.ndim == NUM_MULTI_CHANNEL_DIMENSIONS:
+        transformed_masks = transformed_masks.transpose(2, 0, 1)
+
+    # Normalize the returned bboxes
+    bboxes[:, :4] = bboxes_from_masks(transformed_masks)
+
+    return bboxes
+
+
+@handle_empty_array("keypoints")
+def distort_image_keypoints(
+    keypoints: np.ndarray,
+    generated_mesh: np.ndarray,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Distort keypoints based on a generated mesh.
+
+    This function applies a perspective transformation to each keypoint based on the provided generated mesh.
+    It ensures that the keypoints are clipped to the image boundaries after transformation.
+
+    Args:
+        keypoints (np.ndarray): The keypoints to distort.
+        generated_mesh (np.ndarray): The generated mesh to distort the keypoints with.
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: The distorted keypoints.
+
+    """
+    distorted_keypoints = keypoints.copy()
+    height, width = image_shape[:2]
+
+    for mesh in generated_mesh:
+        x1, y1, x2, y2 = mesh[:4]  # Source rectangle
+        dst_quad = mesh[4:].reshape(4, 2)  # Destination quadrilateral
+
+        src_quad = np.array(
+            [
+                [x1, y1],  # Top-left
+                [x2, y1],  # Top-right
+                [x2, y2],  # Bottom-right
+                [x1, y2],  # Bottom-left
+            ],
+            dtype=np.float32,
+        )
+
+        perspective_mat = cv2.getPerspectiveTransform(src_quad, dst_quad)
+
+        mask = (keypoints[:, 0] >= x1) & (keypoints[:, 0] < x2) & (keypoints[:, 1] >= y1) & (keypoints[:, 1] < y2)
+        cell_keypoints = keypoints[mask]
+
+        if len(cell_keypoints) > 0:
+            # Convert to float32 before applying the transformation
+            points_float32 = cell_keypoints[:, :2].astype(np.float32).reshape(-1, 1, 2)
+            transformed_points = cv2.perspectiveTransform(
+                points_float32,
+                perspective_mat,
+            ).reshape(-1, 2)
+
+            # Update distorted keypoints
+            distorted_keypoints[mask, :2] = transformed_points
+
+    # Clip keypoints to image boundaries
+    distorted_keypoints[:, 0] = np.clip(
+        distorted_keypoints[:, 0],
+        0,
+        width - 1,
+        out=distorted_keypoints[:, 0],
+    )
+    distorted_keypoints[:, 1] = np.clip(
+        distorted_keypoints[:, 1],
+        0,
+        height - 1,
+        out=distorted_keypoints[:, 1],
+    )
+
+    return distorted_keypoints
+
+
+def generate_distorted_grid_polygons(
+    dimensions: np.ndarray,
+    magnitude: int,
+    random_generator: np.random.Generator,
+) -> np.ndarray:
+    """Generate distorted grid polygons based on input dimensions and magnitude.
+
+    This function creates a grid of polygons and applies random distortions to the internal vertices,
+    while keeping the boundary vertices fixed. The distortion is applied consistently across shared
+    vertices to avoid gaps or overlaps in the resulting grid.
+
+    Args:
+        dimensions (np.ndarray): A 3D array of shape (grid_height, grid_width, 4) where each element
+                                 is [x_min, y_min, x_max, y_max] representing the dimensions of a grid cell.
+        magnitude (int): Maximum pixel-wise displacement for distortion. The actual displacement
+                         will be randomly chosen in the range [-magnitude, magnitude].
+        random_generator (np.random.Generator): A random number generator.
+
+    Returns:
+        np.ndarray: A 2D array of shape (total_cells, 8) where each row represents a distorted polygon
+                    as [x1, y1, x2, y1, x2, y2, x1, y2]. The total_cells is equal to grid_height * grid_width.
+
+    Note:
+        - Only internal grid points are distorted; boundary points remain fixed.
+        - The function ensures consistent distortion across shared vertices of adjacent cells.
+        - The distortion is applied to the following points of each internal cell:
+            * Bottom-right of the cell above and to the left
+            * Bottom-left of the cell above
+            * Top-right of the cell to the left
+            * Top-left of the current cell
+        - Each square represents a cell, and the X marks indicate the coordinates where displacement occurs.
+            +--+--+--+--+
+            |  |  |  |  |
+            +--X--X--X--+
+            |  |  |  |  |
+            +--X--X--X--+
+            |  |  |  |  |
+            +--X--X--X--+
+            |  |  |  |  |
+            +--+--+--+--+
+        - For each X, the coordinates of the left, right, top, and bottom edges
+          in the four adjacent cells are displaced.
+
+    Examples:
+        >>> dimensions = np.array([[[0, 0, 50, 50], [50, 0, 100, 50]],
+        ...                        [[0, 50, 50, 100], [50, 50, 100, 100]]])
+        >>> distorted = generate_distorted_grid_polygons(dimensions, magnitude=10)
+        >>> distorted.shape
+        (4, 8)
+
+    """
+    grid_height, grid_width = dimensions.shape[:2]
+    total_cells = grid_height * grid_width
+
+    # Initialize polygons
+    polygons = np.zeros((total_cells, 8), dtype=np.float32)
+    polygons[:, 0:2] = dimensions.reshape(-1, 4)[:, [0, 1]]  # x1, y1
+    polygons[:, 2:4] = dimensions.reshape(-1, 4)[:, [2, 1]]  # x2, y1
+    polygons[:, 4:6] = dimensions.reshape(-1, 4)[:, [2, 3]]  # x2, y2
+    polygons[:, 6:8] = dimensions.reshape(-1, 4)[:, [0, 3]]  # x1, y2
+
+    # Generate displacements for internal grid points only
+    internal_points_height, internal_points_width = grid_height - 1, grid_width - 1
+    displacements = random_generator.integers(
+        -magnitude,
+        magnitude + 1,
+        size=(internal_points_height, internal_points_width, 2),
+    ).astype(np.float32)
+
+    # Apply displacements to internal polygon vertices
+    for i in range(1, grid_height):
+        for j in range(1, grid_width):
+            dx, dy = displacements[i - 1, j - 1]
+
+            # Bottom-right of cell (i-1, j-1)
+            polygons[(i - 1) * grid_width + (j - 1), 4:6] += [dx, dy]
+
+            # Bottom-left of cell (i-1, j)
+            polygons[(i - 1) * grid_width + j, 6:8] += [dx, dy]
+
+            # Top-right of cell (i, j-1)
+            polygons[i * grid_width + (j - 1), 2:4] += [dx, dy]
+
+            # Top-left of cell (i, j)
+            polygons[i * grid_width + j, 0:2] += [dx, dy]
+
+    return polygons
+
+
+@handle_empty_array("keypoints")
+def pad_keypoints(
+    keypoints: np.ndarray,
+    pad_top: int,
+    pad_bottom: int,
+    pad_left: int,
+    pad_right: int,
+    border_mode: int,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Pad keypoints by a given amount.
+
+    This function pads keypoints by a given amount.
+    It handles both reflection and padding.
+
+    Args:
+        keypoints (np.ndarray): The keypoints to pad.
+        pad_top (int): The amount to pad the top of the keypoints.
+        pad_bottom (int): The amount to pad the bottom of the keypoints.
+        pad_left (int): The amount to pad the left of the keypoints.
+        pad_right (int): The amount to pad the right of the keypoints.
+        border_mode (int): The border mode to use.
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: The padded keypoints.
+
+    """
+    if border_mode not in REFLECT_BORDER_MODES:
+        shift_vector = np.array([pad_left, pad_top, 0])
+        return shift_keypoints(keypoints, shift_vector)
+
+    grid_dimensions = get_pad_grid_dimensions(
+        pad_top,
+        pad_bottom,
+        pad_left,
+        pad_right,
+        image_shape,
+    )
+
+    keypoints = generate_reflected_keypoints(keypoints, grid_dimensions, image_shape)
+
+    rows, cols = image_shape[:2]
+
+    # Calculate the number of grid cells added on each side
+    original_row, original_col = grid_dimensions["original_position"]
+
+    # Subtract the offset based on the number of added grid cells
+    keypoints[:, 0] -= original_col * cols - pad_left  # x
+    keypoints[:, 1] -= original_row * rows - pad_top  # y
+
+    new_height = pad_top + pad_bottom + rows
+    new_width = pad_left + pad_right + cols
+
+    return validate_keypoints(keypoints, (new_height, new_width))
+
+
+def validate_keypoints(
+    keypoints: np.ndarray,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Validate keypoints and remove those that fall outside the image boundaries.
+
+    Args:
+        keypoints (np.ndarray): Array of keypoints with shape (N, M) where N is the number of keypoints
+                                and M >= 2. The first two columns represent x and y coordinates.
+        image_shape (tuple[int, int]): Shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: Array of valid keypoints that fall within the image boundaries.
+
+    Note:
+        This function only checks the x and y coordinates (first two columns) of the keypoints.
+        Any additional columns (e.g., angle, scale) are preserved for valid keypoints.
+
+    """
+    rows, cols = image_shape[:2]
+
+    x, y = keypoints[:, 0], keypoints[:, 1]
+
+    valid_indices = (x >= 0) & (x < cols) & (y >= 0) & (y < rows)
+
+    return keypoints[valid_indices]
+
+
+def shift_keypoints(keypoints: np.ndarray, shift_vector: np.ndarray) -> np.ndarray:
+    """Shift keypoints by a given shift vector.
+
+    This function shifts the keypoints by a given shift vector.
+    It only shifts the x, y and z coordinates of the keypoints.
+
+    Args:
+        keypoints (np.ndarray): The keypoints to shift.
+        shift_vector (np.ndarray): The shift vector to apply to the keypoints.
+
+    Returns:
+        np.ndarray: The shifted keypoints.
+
+    """
+    shifted_keypoints = keypoints.copy()
+    shifted_keypoints[:, :3] += shift_vector[:3]  # Only shift x, y and z
+    return shifted_keypoints
+
+
+def generate_reflected_keypoints(
+    keypoints: np.ndarray,
+    grid_dims: dict[str, tuple[int, int]],
+    image_shape: tuple[int, int],
+    center_in_origin: bool = False,
+) -> np.ndarray:
+    """Generate reflected keypoints for the entire reflection grid.
+
+    This function creates a grid of keypoints by reflecting and shifting the original keypoints.
+    It handles both centered and non-centered grids based on the `center_in_origin` parameter.
+
+    Args:
+        keypoints (np.ndarray): Original keypoints array of shape (N, 4+), where N is the number of keypoints,
+                                and each keypoint is represented by at least 4 values (x, y, angle, scale, ...).
+        grid_dims (dict[str, tuple[int, int]]): A dictionary containing grid dimensions and original position.
+            It should have the following keys:
+            - "grid_shape": tuple[int, int] representing (grid_rows, grid_cols)
+            - "original_position": tuple[int, int] representing (original_row, original_col)
+        image_shape (tuple[int, int]): Shape of the original image as (height, width).
+        center_in_origin (bool, optional): If True, center the grid at the origin. Default is False.
+
+    Returns:
+        np.ndarray: Array of reflected and shifted keypoints for the entire grid. The shape is
+                    (N * grid_rows * grid_cols, 4+), where N is the number of original keypoints.
+
+    Note:
+        - The function handles keypoint flipping and shifting to create a grid of reflected keypoints.
+        - It preserves the angle and scale information of the keypoints during transformations.
+        - The resulting grid can be either centered at the origin or positioned based on the original grid.
+
+    """
+    grid_rows, grid_cols = grid_dims["grid_shape"]
+    original_row, original_col = grid_dims["original_position"]
+
+    # Prepare flipped versions of keypoints
+    keypoints_hflipped = flip_keypoints(
+        keypoints,
+        flip_horizontal=True,
+        image_shape=image_shape,
+    )
+    keypoints_vflipped = flip_keypoints(
+        keypoints,
+        flip_vertical=True,
+        image_shape=image_shape,
+    )
+    keypoints_hvflipped = flip_keypoints(
+        keypoints,
+        flip_horizontal=True,
+        flip_vertical=True,
+        image_shape=image_shape,
+    )
+
+    rows, cols = image_shape[:2]
+
+    # Shift all versions to the original position
+    shift_vector = np.array(
+        [original_col * cols, original_row * rows, 0, 0, 0],
+    )  # Only shift x and y
+    keypoints = shift_keypoints(keypoints, shift_vector)
+    keypoints_hflipped = shift_keypoints(keypoints_hflipped, shift_vector)
+    keypoints_vflipped = shift_keypoints(keypoints_vflipped, shift_vector)
+    keypoints_hvflipped = shift_keypoints(keypoints_hvflipped, shift_vector)
+
+    new_keypoints = []
+
+    for grid_row in range(grid_rows):
+        for grid_col in range(grid_cols):
+            # Determine which version of keypoints to use based on grid position
+            if (grid_row - original_row) % 2 == 0 and (grid_col - original_col) % 2 == 0:
+                current_keypoints = keypoints
+            elif (grid_row - original_row) % 2 == 0:
+                current_keypoints = keypoints_hflipped
+            elif (grid_col - original_col) % 2 == 0:
+                current_keypoints = keypoints_vflipped
+            else:
+                current_keypoints = keypoints_hvflipped
+
+            # Shift to the current grid cell
+            cell_shift = np.array(
+                [
+                    (grid_col - original_col) * cols,
+                    (grid_row - original_row) * rows,
+                    0,
+                    0,
+                    0,
+                ],
+            )
+            shifted_keypoints = shift_keypoints(current_keypoints, cell_shift)
+
+            new_keypoints.append(shifted_keypoints)
+
+    result = np.vstack(new_keypoints)
+
+    return shift_keypoints(result, -shift_vector) if center_in_origin else result
+
+
+@handle_empty_array("keypoints")
+def flip_keypoints(
+    keypoints: np.ndarray,
+    flip_horizontal: bool = False,
+    flip_vertical: bool = False,
+    image_shape: tuple[int, int] = (0, 0),
+) -> np.ndarray:
+    """Flip keypoints horizontally or vertically.
+
+    This function flips keypoints horizontally or vertically based on the provided parameters.
+    It also flips the angle of the keypoints when flipping horizontally.
+
+    Args:
+        keypoints (np.ndarray): The keypoints to flip.
+        flip_horizontal (bool): Whether to flip horizontally.
+        flip_vertical (bool): Whether to flip vertically.
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+
+    Returns:
+        np.ndarray: The flipped keypoints.
+
+    """
+    rows, cols = image_shape[:2]
+    flipped_keypoints = keypoints.copy()
+    if flip_horizontal:
+        flipped_keypoints[:, 0] = cols - flipped_keypoints[:, 0]
+        flipped_keypoints[:, 3] = -flipped_keypoints[:, 3]  # Flip angle
+    if flip_vertical:
+        flipped_keypoints[:, 1] = rows - flipped_keypoints[:, 1]
+        flipped_keypoints[:, 3] = -flipped_keypoints[:, 3]  # Flip angle
+    return flipped_keypoints
+
+
+def create_affine_transformation_matrix(
+    translate: Mapping[str, float],
+    shear: dict[str, float],
+    scale: dict[str, float],
+    rotate: float,
+    shift: tuple[float, float],
+) -> np.ndarray:
+    """Create an affine transformation matrix combining translation, shear, scale, and rotation.
+
+    Args:
+        translate (dict[str, float]): Translation in x and y directions.
+        shear (dict[str, float]): Shear in x and y directions (in degrees).
+        scale (dict[str, float]): Scale factors for x and y directions.
+        rotate (float): Rotation angle in degrees.
+        shift (tuple[float, float]): Shift to apply before and after transformations.
+
+    Returns:
+        np.ndarray: The resulting 3x3 affine transformation matrix.
+
+    """
+    # Convert angles to radians
+    rotate_rad = np.deg2rad(rotate % 360)
+
+    shear_x_rad = np.deg2rad(shear["x"])
+    shear_y_rad = np.deg2rad(shear["y"])
+
+    # Create individual transformation matrices
+    # 1. Shift to top-left
+    m_shift_topleft = np.array([[1, 0, -shift[0]], [0, 1, -shift[1]], [0, 0, 1]])
+
+    # 2. Scale
+    m_scale = np.array([[scale["x"], 0, 0], [0, scale["y"], 0], [0, 0, 1]])
+
+    # 3. Rotation
+    m_rotate = np.array(
+        [
+            [np.cos(rotate_rad), np.sin(rotate_rad), 0],
+            [-np.sin(rotate_rad), np.cos(rotate_rad), 0],
+            [0, 0, 1],
+        ],
+    )
+
+    # 4. Shear
+    m_shear = np.array(
+        [[1, np.tan(shear_x_rad), 0], [np.tan(shear_y_rad), 1, 0], [0, 0, 1]],
+    )
+
+    # 5. Translation
+    m_translate = np.array([[1, 0, translate["x"]], [0, 1, translate["y"]], [0, 0, 1]])
+
+    # 6. Shift back to center
+    m_shift_center = np.array([[1, 0, shift[0]], [0, 1, shift[1]], [0, 0, 1]])
+
+    # Combine all transformations
+    # The order is important: transformations are applied from right to left
+    m = m_shift_center @ m_translate @ m_shear @ m_rotate @ m_scale @ m_shift_topleft
+
+    # Ensure the last row is exactly [0, 0, 1]
+    m[2] = [0, 0, 1]
+
+    return m
+
+
+def compute_transformed_image_bounds(
+    matrix: np.ndarray,
+    image_shape: tuple[int, int],
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute the bounds of an image after applying an affine transformation.
+
+    Args:
+        matrix (np.ndarray): The 3x3 affine transformation matrix.
+        image_shape (Tuple[int, int]): The shape of the image as (height, width).
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: A tuple containing:
+            - min_coords: An array with the minimum x and y coordinates.
+            - max_coords: An array with the maximum x and y coordinates.
+
+    """
+    height, width = image_shape[:2]
+
+    # Define the corners of the image
+    corners = np.array([[0, 0, 1], [width, 0, 1], [width, height, 1], [0, height, 1]])
+
+    # Transform the corners
+    transformed_corners = corners @ matrix.T
+    transformed_corners = transformed_corners[:, :2] / transformed_corners[:, 2:]
+
+    # Calculate the bounding box of the transformed corners
+    min_coords = np.floor(transformed_corners.min(axis=0)).astype(int)
+    max_coords = np.ceil(transformed_corners.max(axis=0)).astype(int)
+
+    return min_coords, max_coords
+
+
+def compute_affine_warp_output_shape(
+    matrix: np.ndarray,
+    input_shape: tuple[int, ...],
+) -> tuple[np.ndarray, tuple[int, int]]:
+    """Compute the output shape of an affine warp.
+
+    This function computes the output shape of an affine warp based on the input matrix and input shape.
+    It calculates the transformed image bounds and then determines the output shape based on the input shape.
+
+    Args:
+        matrix (np.ndarray): The 3x3 affine transformation matrix.
+        input_shape (tuple[int, ...]): The shape of the input image as (height, width, ...).
+
+    Returns:
+        tuple[np.ndarray, tuple[int, int]]: A tuple containing:
+            - matrix: The 3x3 affine transformation matrix.
+            - output_shape: The output shape of the affine warp.
+
+    """
+    height, width = input_shape[:2]
+
+    if height == 0 or width == 0:
+        return matrix, cast("tuple[int, int]", input_shape[:2])
+
+    min_coords, max_coords = compute_transformed_image_bounds(matrix, (height, width))
+    minc, minr = min_coords
+    maxc, maxr = max_coords
+
+    out_height = maxr - minr + 1
+    out_width = maxc - minc + 1
+
+    if len(input_shape) == NUM_MULTI_CHANNEL_DIMENSIONS:
+        output_shape = np.ceil((out_height, out_width, input_shape[2]))
+    else:
+        output_shape = np.ceil((out_height, out_width))
+
+    output_shape_tuple = tuple(int(v) for v in output_shape.tolist())
+
+    # fit output image in new shape
+    translation = np.array([[1, 0, -minc], [0, 1, -minr], [0, 0, 1]])
+    matrix = translation @ matrix
+
+    return matrix, cast("tuple[int, int]", output_shape_tuple)
+
+
+def center(image_shape: tuple[int, int]) -> tuple[float, float]:
+    """Calculate the center coordinates if image. Used by images, masks and keypoints.
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image.
+
+    Returns:
+        tuple[float, float]: center_x, center_y
+
+    """
+    height, width = image_shape[:2]
+    return width / 2 - 0.5, height / 2 - 0.5
+
+
+def center_bbox(image_shape: tuple[int, int]) -> tuple[float, float]:
+    """Calculate the center coordinates for of image for bounding boxes.
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image.
+
+    Returns:
+        tuple[float, float]: center_x, center_y
+
+    """
+    height, width = image_shape[:2]
+    return width / 2, height / 2
+
+
+def generate_grid(
+    image_shape: tuple[int, int],
+    steps_x: list[float],
+    steps_y: list[float],
+    num_steps: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate a distorted grid for image transformation based on given step sizes.
+
+    This function creates two 2D arrays (map_x and map_y) that represent a distorted version
+    of the original image grid. These arrays can be used with OpenCV's remap function to
+    apply grid distortion to an image.
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+        steps_x (list[float]): List of step sizes for the x-axis distortion. The length
+            should be num_steps + 1. Each value represents the relative step size for
+            a segment of the grid in the x direction.
+        steps_y (list[float]): List of step sizes for the y-axis distortion. The length
+            should be num_steps + 1. Each value represents the relative step size for
+            a segment of the grid in the y direction.
+        num_steps (int): The number of steps to divide each axis into. This determines
+            the granularity of the distortion grid.
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: A tuple containing two 2D numpy arrays:
+            - map_x: A 2D array of float32 values representing the x-coordinates
+              of the distorted grid.
+            - map_y: A 2D array of float32 values representing the y-coordinates
+              of the distorted grid.
+
+    Note:
+        - The function generates a grid where each cell can be distorted independently.
+        - The distortion is controlled by the steps_x and steps_y parameters, which
+          determine how much each grid line is shifted.
+        - The resulting map_x and map_y can be used directly with cv2.remap() to
+          apply the distortion to an image.
+        - The distortion is applied smoothly across each grid cell using linear
+          interpolation.
+
+    Examples:
+        >>> image_shape = (100, 100)
+        >>> steps_x = [1.1, 0.9, 1.0, 1.2, 0.95, 1.05]
+        >>> steps_y = [0.9, 1.1, 1.0, 1.1, 0.9, 1.0]
+        >>> num_steps = 5
+        >>> map_x, map_y = generate_grid(image_shape, steps_x, steps_y, num_steps)
+        >>> distorted_image = cv2.remap(image, map_x, map_y, cv2.INTER_LINEAR)
+
+    """
+    height, width = image_shape[:2]
+    x_step = width // num_steps
+    xx = np.zeros(width, np.float32)
+    prev = 0.0
+    for idx, step in enumerate(steps_x):
+        x = idx * x_step
+        start = int(x)
+        end = min(int(x) + x_step, width)
+        cur = prev + x_step * step
+        xx[start:end] = np.linspace(prev, cur, end - start)
+        prev = cur
+
+    y_step = height // num_steps
+    yy = np.zeros(height, np.float32)
+    prev = 0.0
+    for idx, step in enumerate(steps_y):
+        y = idx * y_step
+        start = int(y)
+        end = min(int(y) + y_step, height)
+        cur = prev + y_step * step
+        yy[start:end] = np.linspace(prev, cur, end - start)
+        prev = cur
+
+    return np.meshgrid(xx, yy)
+
+
+def normalize_grid_distortion_steps(
+    image_shape: tuple[int, int],
+    num_steps: int,
+    x_steps: list[float],
+    y_steps: list[float],
+) -> dict[str, np.ndarray]:
+    """Normalize the grid distortion steps.
+
+    This function normalizes the grid distortion steps, ensuring that the distortion never leaves the image bounds.
+    It compensates for smaller last steps in the source image and normalizes the steps such that the distortion
+    never leaves the image bounds.
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+        num_steps (int): The number of steps to divide each axis into. This determines
+            the granularity of the distortion grid.
+        x_steps (list[float]): List of step sizes for the x-axis distortion. The length
+            should be num_steps + 1. Each value represents the relative step size for
+            a segment of the grid in the x direction.
+        y_steps (list[float]): List of step sizes for the y-axis distortion. The length
+            should be num_steps + 1. Each value represents the relative step size for
+            a segment of the grid in the y direction.
+
+    Returns:
+        dict[str, np.ndarray]: A dictionary containing the normalized step sizes for the x and y axes.
+
+    """
+    height, width = image_shape[:2]
+
+    # compensate for smaller last steps in source image.
+    x_step = width // num_steps
+    last_x_step = min(width, ((num_steps + 1) * x_step)) - (num_steps * x_step)
+    x_steps[-1] *= last_x_step / x_step
+
+    y_step = height // num_steps
+    last_y_step = min(height, ((num_steps + 1) * y_step)) - (num_steps * y_step)
+    y_steps[-1] *= last_y_step / y_step
+
+    # now normalize such that distortion never leaves image bounds.
+    tx = width / math.floor(width / num_steps)
+    ty = height / math.floor(height / num_steps)
+    x_steps = np.array(x_steps) * (tx / np.sum(x_steps))
+    y_steps = np.array(y_steps) * (ty / np.sum(y_steps))
+
+    return {"steps_x": x_steps, "steps_y": y_steps}
+
+
+def almost_equal_intervals(n: int, parts: int) -> np.ndarray:
+    """Generates an array of nearly equal integer intervals that sum up to `n`.
+
+    This function divides the number `n` into `parts` nearly equal parts. It ensures that
+    the sum of all parts equals `n`, and the difference between any two parts is at most one.
+    This is useful for distributing a total amount into nearly equal discrete parts.
+
+    Args:
+        n (int): The total value to be split.
+        parts (int): The number of parts to split into.
+
+    Returns:
+        np.ndarray: An array of integers where each integer represents the size of a part.
+
+    Examples:
+        >>> almost_equal_intervals(20, 3)
+        array([7, 7, 6])  # Splits 20 into three parts: 7, 7, and 6
+        >>> almost_equal_intervals(16, 4)
+        array([4, 4, 4, 4])  # Splits 16 into four equal parts
+
+    """
+    part_size, remainder = divmod(n, parts)
+    # Create an array with the base part size and adjust the first `remainder` parts by adding 1
+    return np.array(
+        [part_size + 1 if i < remainder else part_size for i in range(parts)],
+    )
+
+
+def generate_shuffled_splits(
+    size: int,
+    divisions: int,
+    random_generator: np.random.Generator,
+) -> np.ndarray:
+    """Generate shuffled splits for a given dimension size and number of divisions.
+
+    Args:
+        size (int): Total size of the dimension (height or width).
+        divisions (int): Number of divisions (rows or columns).
+        random_generator (np.random.Generator | None): The random generator to use for shuffling the splits.
+            If None, the splits are not shuffled.
+
+    Returns:
+        np.ndarray: Cumulative edges of the shuffled intervals.
+
+    """
+    intervals = almost_equal_intervals(size, divisions)
+    random_generator.shuffle(intervals)
+    return np.insert(np.cumsum(intervals), 0, 0)
+
+
+def split_uniform_grid(
+    image_shape: tuple[int, int],
+    grid: tuple[int, int],
+    random_generator: np.random.Generator,
+) -> np.ndarray:
+    """Splits an image shape into a uniform grid specified by the grid dimensions.
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+        grid (tuple[int, int]): The grid size as (rows, columns).
+        random_generator (np.random.Generator): The random generator to use for shuffling the splits.
+            If None, the splits are not shuffled.
+
+    Returns:
+        np.ndarray: An array containing the tiles' coordinates in the format (start_y, start_x, end_y, end_x).
+
+    Note:
+        The function uses `generate_shuffled_splits` to generate the splits for the height and width of the image.
+        The splits are then used to calculate the coordinates of the tiles.
+
+    """
+    n_rows, n_cols = grid
+
+    height_splits = generate_shuffled_splits(
+        image_shape[0],
+        grid[0],
+        random_generator=random_generator,
+    )
+    width_splits = generate_shuffled_splits(
+        image_shape[1],
+        grid[1],
+        random_generator=random_generator,
+    )
+
+    # Calculate tiles coordinates
+    tiles = [
+        (height_splits[i], width_splits[j], height_splits[i + 1], width_splits[j + 1])
+        for i in range(n_rows)
+        for j in range(n_cols)
+    ]
+
+    return np.array(tiles, dtype=np.int16)
+
+
+def generate_perspective_points(
+    image_shape: tuple[int, int],
+    scale: float,
+    random_generator: np.random.Generator,
+) -> np.ndarray:
+    """Generate perspective points for a given image shape and scale.
+
+    This function generates perspective points for a given image shape and scale.
+    It uses a normal distribution to generate the points, and then modulates them to be within the image bounds.
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+        scale (float): The scale of the perspective points.
+        random_generator (np.random.Generator): The random generator to use for generating the points.
+
+    Returns:
+        np.ndarray: The perspective points.
+
+    """
+    height, width = image_shape[:2]
+    points = random_generator.normal(0, scale, (4, 2))
+    points = np.mod(np.abs(points), 0.32)
+
+    # top left -- no changes needed, just use jitter
+    # top right
+    points[1, 0] = 1.0 - points[1, 0]  # w = 1.0 - jitter
+    # bottom right
+    points[2] = 1.0 - points[2]  # w = 1.0 - jitter
+    # bottom left
+    points[3, 1] = 1.0 - points[3, 1]  # h = 1.0 - jitter
+
+    points[:, 0] *= width
+    points[:, 1] *= height
+
+    return points
+
+
+def order_points(pts: np.ndarray) -> np.ndarray:
+    """Order points in a clockwise manner.
+
+    This function orders the points in a clockwise manner, ensuring that the points are in the correct
+    order for perspective transformation.
+
+    Args:
+        pts (np.ndarray): The points to order.
+
+    Returns:
+        np.ndarray: The ordered points.
+
+    """
+    pts = np.array(sorted(pts, key=lambda x: x[0]))
+    left = pts[:2]  # points with smallest x coordinate - left points
+    right = pts[2:]  # points with greatest x coordinate - right points
+
+    if left[0][1] < left[1][1]:
+        tl, bl = left
+    else:
+        bl, tl = left
+
+    if right[0][1] < right[1][1]:
+        tr, br = right
+    else:
+        br, tr = right
+
+    return np.array([tl, tr, br, bl], dtype=np.float32)
+
+
+def compute_perspective_params(
+    points: np.ndarray,
+    image_shape: tuple[int, int],
+) -> tuple[np.ndarray, int, int]:
+    """Compute perspective transformation parameters.
+
+    This function computes the perspective transformation parameters for a given set of points.
+    It adjusts the points to ensure that the transformed image retains its original dimensions.
+
+    Args:
+        points (np.ndarray): The points to compute the perspective transformation parameters for.
+        image_shape (tuple[int, int]): The shape of the image.
+
+    Returns:
+        tuple[np.ndarray, int, int]: The perspective transformation parameters and the maximum
+            dimensions of the transformed image.
+
+    """
+    height, width = image_shape
+    top_left, top_right, bottom_right, bottom_left = points
+
+    def adjust_dimension(
+        dim1: np.ndarray,
+        dim2: np.ndarray,
+        min_size: int = 2,
+    ) -> float:
+        size = np.sqrt(np.sum((dim1 - dim2) ** 2))
+        if size < min_size:
+            step_size = (min_size - size) / 2
+            dim1[dim1 > dim2] += step_size
+            dim2[dim1 > dim2] -= step_size
+            dim1[dim1 <= dim2] -= step_size
+            dim2[dim1 <= dim2] += step_size
+            size = min_size
+        return size
+
+    max_width = max(
+        adjust_dimension(top_right, top_left),
+        adjust_dimension(bottom_right, bottom_left),
+    )
+    max_height = max(
+        adjust_dimension(bottom_right, top_right),
+        adjust_dimension(bottom_left, top_left),
+    )
+
+    dst = np.array([[0, 0], [width, 0], [width, height], [0, height]], dtype=np.float32)
+    matrix = cv2.getPerspectiveTransform(points, dst)
+
+    return matrix, int(max_width), int(max_height)
+
+
+def expand_transform(
+    matrix: np.ndarray,
+    shape: tuple[int, int],
+) -> tuple[np.ndarray, int, int]:
+    """Expand a transformation matrix to include padding.
+
+    This function expands a transformation matrix to include padding, ensuring that the transformed
+    image retains its original dimensions. It first calculates the destination points of the transformed
+    image, then adjusts the matrix to include padding, and finally returns the expanded matrix and the
+    maximum dimensions of the transformed image.
+
+    Args:
+        matrix (np.ndarray): The transformation matrix to expand.
+        shape (tuple[int, int]): The shape of the image.
+
+    Returns:
+        tuple[np.ndarray, int, int]: The expanded matrix and the maximum dimensions of the transformed image.
+
+    """
+    height, width = shape[:2]
+    rect = np.array(
+        [[0, 0], [width, 0], [width, height], [0, height]],
+        dtype=np.float32,
+    )
+    dst = cv2.perspectiveTransform(np.array([rect]), matrix)[0]
+
+    dst -= dst.min(axis=0, keepdims=True)
+    dst = np.around(dst, decimals=0)
+
+    matrix_expanded = cv2.getPerspectiveTransform(rect, dst)
+    max_width, max_height = dst.max(axis=0)
+    return matrix_expanded, int(max_width), int(max_height)
+
+
+def create_piecewise_affine_maps(
+    image_shape: tuple[int, int],
+    grid: tuple[int, int],
+    scale: float,
+    absolute_scale: bool,
+    random_generator: np.random.Generator,
+) -> tuple[np.ndarray | None, np.ndarray | None]:
+    """Create maps for piecewise affine transformation using OpenCV's remap function.
+
+    This function creates maps for piecewise affine transformation using OpenCV's remap function.
+    It generates the control points for the transformation, then uses the remap function to create
+    the transformation maps.
+
+    Args:
+        image_shape (tuple[int, int]): The shape of the image as (height, width).
+        grid (tuple[int, int]): The grid size as (rows, columns).
+        scale (float): The scale of the transformation.
+        absolute_scale (bool): Whether to use absolute scale.
+        random_generator (np.random.Generator): The random generator to use for generating the points.
+
+    Returns:
+        tuple[np.ndarray | None, np.ndarray | None]: The transformation maps.
+
+    """
+    height, width = image_shape[:2]
+    nb_rows, nb_cols = grid
+
+    # Input validation
+    if height <= 0 or width <= 0 or nb_rows <= 0 or nb_cols <= 0:
+        raise ValueError("Dimensions must be positive")
+    if scale <= 0:
+        return None, None
+
+    # Create source points grid
+    y = np.linspace(0, height - 1, nb_rows, dtype=np.float32)
+    x = np.linspace(0, width - 1, nb_cols, dtype=np.float32)
+    xx_src, yy_src = np.meshgrid(x, y)
+
+    # Initialize destination maps at full resolution
+    map_x = np.zeros((height, width), dtype=np.float32)
+    map_y = np.zeros((height, width), dtype=np.float32)
+
+    # Generate jitter for control points
+    jitter_scale = scale / 3 if absolute_scale else scale * min(width, height) / 3
+
+    jitter = random_generator.normal(0, jitter_scale, (nb_rows, nb_cols, 2)).astype(
+        np.float32,
+    )
+
+    # Create control points with jitter
+    control_points = np.zeros((nb_rows * nb_cols, 4), dtype=np.float32)
+    for i in range(nb_rows):
+        for j in range(nb_cols):
+            idx = i * nb_cols + j
+            # Source points
+            control_points[idx, 0] = xx_src[i, j]
+            control_points[idx, 1] = yy_src[i, j]
+            # Destination points with jitter
+            control_points[idx, 2] = np.clip(
+                xx_src[i, j] + jitter[i, j, 1],
+                0,
+                width - 1,
+            )
+            control_points[idx, 3] = np.clip(
+                yy_src[i, j] + jitter[i, j, 0],
+                0,
+                height - 1,
+            )
+
+    # Create full resolution maps
+    for i in range(height):
+        for j in range(width):
+            # Find nearest control points and interpolate
+            dx = j - control_points[:, 0]
+            dy = i - control_points[:, 1]
+            dist = dx * dx + dy * dy
+            weights = 1 / (dist + 1e-8)
+            weights = weights / np.sum(weights)
+
+            map_x[i, j] = np.sum(weights * control_points[:, 2])
+            map_y[i, j] = np.sum(weights * control_points[:, 3])
+
+    # Ensure output is within bounds
+    map_x = np.clip(map_x, 0, width - 1, out=map_x)
+    map_y = np.clip(map_y, 0, height - 1, out=map_y)
+
+    return map_x, map_y
+
+
+@handle_empty_array("bboxes")
+def bboxes_piecewise_affine(
+    bboxes: np.ndarray,
+    map_x: np.ndarray,
+    map_y: np.ndarray,
+    border_mode: int,
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Apply a piecewise affine transformation to bounding boxes.
+
+    This function applies a piecewise affine transformation to the bounding boxes of an image.
+    It first converts the bounding boxes to masks, then applies the transformation, and finally
+    converts the transformed masks back to bounding boxes.
+
+    Args:
+        bboxes (np.ndarray): The bounding boxes to transform.
+        map_x (np.ndarray): The x-coordinates of the transformation.
+        map_y (np.ndarray): The y-coordinates of the transformation.
+        border_mode (int): The border mode to use for the transformation.
+        image_shape (tuple[int, int]): The shape of the image.
+
+    Returns:
+        np.ndarray: The transformed bounding boxes.
+
+    """
+    masks = masks_from_bboxes(bboxes, image_shape).transpose(1, 2, 0)
+
+    map_xy = np.stack([map_x, map_y], axis=-1).astype(np.float32)
+
+    # Call remap with the combined map and empty second map
+    transformed_masks = cv2.remap(
+        masks,
+        map_xy,
+        None,
+        cv2.INTER_NEAREST,
+        borderMode=border_mode,
+        borderValue=0,
+    )
+
+    if transformed_masks.ndim == NUM_MULTI_CHANNEL_DIMENSIONS:
+        transformed_masks = transformed_masks.transpose(2, 0, 1)
+
+    # Normalize the returned bboxes
+    bboxes[:, :4] = bboxes_from_masks(transformed_masks)
+
+    return bboxes
+
+
+def get_dimension_padding(
+    current_size: int,
+    min_size: int | None,
+    divisor: int | None,
+) -> tuple[int, int]:
+    """Calculate padding for a single dimension.
+
+    Args:
+        current_size (int): Current size of the dimension
+        min_size (int | None): Minimum size requirement, if any
+        divisor (int | None): Divisor for padding to make size divisible, if any
+
+    Returns:
+        tuple[int, int]: (pad_before, pad_after)
+
+    """
+    if min_size is not None:
+        if current_size < min_size:
+            pad_before = int((min_size - current_size) / 2.0)
+            pad_after = min_size - current_size - pad_before
+            return pad_before, pad_after
+    elif divisor is not None:
+        remainder = current_size % divisor
+        if remainder > 0:
+            total_pad = divisor - remainder
+            pad_before = total_pad // 2
+            pad_after = total_pad - pad_before
+            return pad_before, pad_after
+
+    return 0, 0
+
+
+def get_padding_params(
+    image_shape: tuple[int, int],
+    min_height: int | None,
+    min_width: int | None,
+    pad_height_divisor: int | None,
+    pad_width_divisor: int | None,
+) -> tuple[int, int, int, int]:
+    """Calculate padding parameters based on target dimensions.
+
+    Args:
+        image_shape (tuple[int, int]): (height, width) of the image
+        min_height (int | None): Minimum height requirement, if any
+        min_width (int | None): Minimum width requirement, if any
+        pad_height_divisor (int | None): Divisor for height padding, if any
+        pad_width_divisor (int | None): Divisor for width padding, if any
+
+    Returns:
+        tuple[int, int, int, int]: (pad_top, pad_bottom, pad_left, pad_right)
+
+    """
+    rows, cols = image_shape[:2]
+
+    h_pad_top, h_pad_bottom = get_dimension_padding(
+        rows,
+        min_height,
+        pad_height_divisor,
+    )
+    w_pad_left, w_pad_right = get_dimension_padding(cols, min_width, pad_width_divisor)
+
+    return h_pad_top, h_pad_bottom, w_pad_left, w_pad_right
+
+
+def adjust_padding_by_position(
+    h_top: int,
+    h_bottom: int,
+    w_left: int,
+    w_right: int,
+    position: Literal["center", "top_left", "top_right", "bottom_left", "bottom_right", "random"],
+    py_random: np.random.RandomState,
+) -> tuple[int, int, int, int]:
+    """Adjust padding values based on desired position."""
+    if position == "center":
+        return h_top, h_bottom, w_left, w_right
+
+    if position == "top_left":
+        return 0, h_top + h_bottom, 0, w_left + w_right
+
+    if position == "top_right":
+        return 0, h_top + h_bottom, w_left + w_right, 0
+
+    if position == "bottom_left":
+        return h_top + h_bottom, 0, 0, w_left + w_right
+
+    if position == "bottom_right":
+        return h_top + h_bottom, 0, w_left + w_right, 0
+
+    if position == "random":
+        h_pad = h_top + h_bottom
+        w_pad = w_left + w_right
+        h_top = py_random.randint(0, h_pad)
+        h_bottom = h_pad - h_top
+        w_left = py_random.randint(0, w_pad)
+        w_right = w_pad - w_left
+        return h_top, h_bottom, w_left, w_right
+
+    raise ValueError(f"Unknown position: {position}")
+
+
+def swap_tiles_on_keypoints(
+    keypoints: np.ndarray,
+    tiles: np.ndarray,
+    mapping: np.ndarray,
+) -> np.ndarray:
+    """Swap the positions of keypoints based on a tile mapping.
+
+    This function takes a set of keypoints and repositions them according to a mapping of tile swaps.
+    Keypoints are moved from their original tiles to new positions in the swapped tiles.
+
+    Args:
+        keypoints (np.ndarray): A 2D numpy array of shape (N, 2) where N is the number of keypoints.
+                                Each row represents a keypoint's (x, y) coordinates.
+        tiles (np.ndarray): A 2D numpy array of shape (M, 4) where M is the number of tiles.
+                            Each row represents a tile's (start_y, start_x, end_y, end_x) coordinates.
+        mapping (np.ndarray): A 1D numpy array of shape (M,) where M is the number of tiles.
+                              Each element i contains the index of the tile that tile i should be swapped with.
+
+    Returns:
+        np.ndarray: A 2D numpy array of the same shape as the input keypoints, containing the new positions
+                    of the keypoints after the tile swap.
+
+    Raises:
+        RuntimeWarning: If any keypoint is not found within any tile.
+
+    Notes:
+        - Keypoints that do not fall within any tile will remain unchanged.
+        - The function assumes that the tiles do not overlap and cover the entire image space.
+
+    """
+    if not keypoints.size:
+        return keypoints
+
+    # Broadcast keypoints and tiles for vectorized comparison
+    kp_x = keypoints[:, 0][:, np.newaxis]  # Shape: (num_keypoints, 1)
+    kp_y = keypoints[:, 1][:, np.newaxis]  # Shape: (num_keypoints, 1)
+
+    start_y, start_x, end_y, end_x = tiles.T  # Each shape: (num_tiles,)
+
+    # Check if each keypoint is inside each tile
+    in_tile = (kp_y >= start_y) & (kp_y < end_y) & (kp_x >= start_x) & (kp_x < end_x)
+
+    # Find which tile each keypoint belongs to
+    tile_indices = np.argmax(in_tile, axis=1)
+
+    # Check if any keypoint is not in any tile
+    not_in_any_tile = ~np.any(in_tile, axis=1)
+    if np.any(not_in_any_tile):
+        warn(
+            "Some keypoints are not in any tile. They will be returned unchanged. This is unexpected and should be "
+            "investigated.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+
+    # Get the new tile indices
+    new_tile_indices = np.array(mapping)[tile_indices]
+
+    # Calculate the offsets
+    old_start_x = tiles[tile_indices, 1]
+    old_start_y = tiles[tile_indices, 0]
+    new_start_x = tiles[new_tile_indices, 1]
+    new_start_y = tiles[new_tile_indices, 0]
+
+    # Apply the transformation
+    new_keypoints = keypoints.copy()
+    new_keypoints[:, 0] = (keypoints[:, 0] - old_start_x) + new_start_x
+    new_keypoints[:, 1] = (keypoints[:, 1] - old_start_y) + new_start_y
+
+    # Keep original coordinates for keypoints not in any tile
+    new_keypoints[not_in_any_tile] = keypoints[not_in_any_tile]
+
+    return new_keypoints
+
+
+def swap_tiles_on_image(
+    image: np.ndarray,
+    tiles: np.ndarray,
+    mapping: list[int] | None = None,
+) -> np.ndarray:
+    """Swap tiles on the image according to the new format.
+
+    Args:
+        image (np.ndarray): Input image.
+        tiles (np.ndarray): Array of tiles with each tile as [start_y, start_x, end_y, end_x].
+        mapping (list[int] | None): list of new tile indices.
+
+    Returns:
+        np.ndarray: Output image with tiles swapped according to the random shuffle.
+
+    """
+    # If no tiles are provided, return a copy of the original image
+    if tiles.size == 0 or mapping is None:
+        return image.copy()
+
+    # Create a copy of the image to retain original for reference
+    new_image = np.empty_like(image)
+    for num, new_index in enumerate(mapping):
+        start_y, start_x, end_y, end_x = tiles[new_index]
+        start_y_orig, start_x_orig, end_y_orig, end_x_orig = tiles[num]
+        # Assign the corresponding tile from the original image to the new image
+        new_image[start_y:end_y, start_x:end_x] = image[
+            start_y_orig:end_y_orig,
+            start_x_orig:end_x_orig,
+        ]
+
+    return new_image
+
+
+def is_valid_component(
+    component_area: float,
+    original_area: float,
+    min_area: float | None,
+    min_visibility: float | None,
+) -> bool:
+    """Validate if a component meets the minimum requirements."""
+    visibility = component_area / original_area
+    return (min_area is None or component_area >= min_area) and (min_visibility is None or visibility >= min_visibility)
+
+
+@handle_empty_array("bboxes")
+def bboxes_grid_shuffle(
+    bboxes: np.ndarray,
+    tiles: np.ndarray,
+    mapping: list[int],
+    image_shape: tuple[int, int],
+    min_area: float,
+    min_visibility: float,
+) -> np.ndarray:
+    """Shuffle bounding boxes according to grid mapping.
+
+    Args:
+        bboxes (np.ndarray): Array of bounding boxes with shape (num_boxes, 4+)
+        tiles (np.ndarray): Array of grid tiles
+        mapping (list[int]): Mapping of tile indices
+        image_shape (tuple[int, int]): Shape of the image (height, width)
+        min_area (float): Minimum area of a bounding box to keep
+        min_visibility (float): Minimum visibility ratio of a bounding box to keep
+
+    Returns:
+        np.ndarray: Shuffled bounding boxes
+
+    """
+    # Convert bboxes to masks
+    masks = masks_from_bboxes(bboxes, image_shape)
+
+    # Apply grid shuffle to each mask and handle split components
+    all_component_masks = []
+    extra_bbox_data = []  # Store additional bbox data for each component
+
+    for idx, mask in enumerate(masks):
+        original_area = np.sum(mask)  # Get original mask area
+
+        # Shuffle the mask
+        shuffled_mask = swap_tiles_on_image(mask, tiles, mapping)
+
+        # Find connected components
+        num_components, components = cv2.connectedComponents(
+            shuffled_mask.astype(np.uint8),
+        )
+
+        # For each component, create a separate binary mask
+        for comp_idx in range(1, num_components):  # Skip background (0)
+            component_mask = (components == comp_idx).astype(np.uint8)
+
+            # Calculate area and visibility ratio
+            component_area = np.sum(component_mask)
+            # Check if component meets minimum requirements
+            if is_valid_component(
+                component_area,
+                original_area,
+                min_area,
+                min_visibility,
+            ):
+                all_component_masks.append(component_mask)
+                # Append additional bbox data for this component
+                if bboxes.shape[1] > NUM_BBOXES_COLUMNS_IN_ALBUMENTATIONS:
+                    extra_bbox_data.append(bboxes[idx, 4:])
+
+    # Convert all component masks to bboxes
+    if all_component_masks:
+        all_component_masks = np.array(all_component_masks)
+        shuffled_bboxes = bboxes_from_masks(all_component_masks)
+
+        # Add back additional bbox data if present
+        if extra_bbox_data:
+            extra_bbox_data = np.array(extra_bbox_data)
+            return np.column_stack([shuffled_bboxes, extra_bbox_data])
+    else:
+        # Handle case where no valid components were found
+        return np.zeros((0, bboxes.shape[1]), dtype=bboxes.dtype)
+
+    return shuffled_bboxes
+
+
+def create_shape_groups(tiles: np.ndarray) -> dict[tuple[int, int], list[int]]:
+    """Groups tiles by their shape and stores the indices for each shape."""
+    shape_groups = defaultdict(list)
+    for index, (start_y, start_x, end_y, end_x) in enumerate(tiles):
+        shape = (end_y - start_y, end_x - start_x)
+        shape_groups[shape].append(index)
+    return shape_groups
+
+
+def shuffle_tiles_within_shape_groups(
+    shape_groups: dict[tuple[int, int], list[int]],
+    random_generator: np.random.Generator,
+) -> list[int]:
+    """Shuffles indices within each group of similar shapes and creates a list where each
+    index points to the index of the tile it should be mapped to.
+
+    Args:
+        shape_groups (dict[tuple[int, int], list[int]]): Groups of tile indices categorized by shape.
+        random_generator (np.random.Generator): The random generator to use for shuffling the indices.
+            If None, a new random generator will be used.
+
+    Returns:
+        list[int]: A list where each index is mapped to the new index of the tile after shuffling.
+
+    """
+    # Initialize the output list with the same size as the total number of tiles, filled with -1
+    num_tiles = sum(len(indices) for indices in shape_groups.values())
+    mapping = [-1] * num_tiles
+
+    # Prepare the random number generator
+
+    for indices in shape_groups.values():
+        shuffled_indices = indices.copy()
+        random_generator.shuffle(shuffled_indices)
+
+        for old, new in zip(indices, shuffled_indices):
+            mapping[old] = new
+
+    return mapping
+
+
+def compute_pairwise_distances(
+    points1: np.ndarray,
+    points2: np.ndarray,
+) -> np.ndarray:
+    """Compute pairwise distances between two sets of points.
+
+    Args:
+        points1 (np.ndarray): First set of points with shape (N, 2)
+        points2 (np.ndarray): Second set of points with shape (M, 2)
+
+    Returns:
+        np.ndarray: Matrix of pairwise distances with shape (N, M)
+
+    """
+    points1 = np.ascontiguousarray(points1, dtype=np.float32)
+    points2 = np.ascontiguousarray(points2, dtype=np.float32)
+
+    # Compute squared terms
+    p1_squared = cv2.multiply(points1, points1).sum(axis=1, keepdims=True)
+    p2_squared = cv2.multiply(points2, points2).sum(axis=1)[None, :]
+
+    # Compute dot product
+    dot_product = cv2.gemm(points1, points2.T, 1, None, 0)
+
+    return p1_squared + p2_squared - 2 * dot_product
+
+
+def compute_tps_weights(
+    src_points: np.ndarray,
+    dst_points: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute Thin Plate Spline weights.
+
+    Args:
+        src_points (np.ndarray): Source control points with shape (num_points, 2)
+        dst_points (np.ndarray): Destination control points with shape (num_points, 2)
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: Tuple of (nonlinear_weights, affine_weights)
+        - nonlinear_weights: TPS kernel weights for nonlinear deformation (num_points, 2)
+        - affine_weights: Weights for affine transformation (3, 2)
+            [constant term, x scale/shear, y scale/shear]
+
+    Note:
+        The TPS interpolation is decomposed into:
+        1. Nonlinear part (controlled by kernel weights)
+        2. Affine part (global scaling, rotation, translation)
+
+    """
+    num_points = src_points.shape[0]
+
+    # Compute pairwise distances
+    distances = compute_pairwise_distances(src_points, src_points)
+
+    kernel_matrix = np.where(
+        distances > 0,
+        distances * distances * cv2.log(distances + 1e-6),
+        0,
+    ).astype(np.float32)
+
+    # Build system matrix efficiently
+    affine_terms = np.empty((num_points, 3), dtype=np.float32)
+    affine_terms[:, 0] = 1
+    affine_terms[:, 1:] = src_points
+
+    # Construct system matrix
+    system_matrix = np.zeros((num_points + 3, num_points + 3), dtype=np.float32)
+    system_matrix[:num_points, :num_points] = kernel_matrix
+    system_matrix[:num_points, num_points:] = affine_terms
+    system_matrix[num_points:, :num_points] = affine_terms.T
+
+    # Prepare target coordinates
+    target = np.zeros((num_points + 3, 2), dtype=np.float32)
+    target[:num_points] = dst_points
+
+    weights = cv2.solve(system_matrix, target, flags=cv2.DECOMP_LU)[1]
+
+    return weights[:num_points], weights[num_points:]
+
+
+def tps_transform(
+    target_points: np.ndarray,
+    control_points: np.ndarray,
+    nonlinear_weights: np.ndarray,
+    affine_weights: np.ndarray,
+) -> np.ndarray:
+    """Apply TPS transformation with consistent types."""
+    # Ensure float32 type for all inputs
+    target_points = np.ascontiguousarray(target_points, dtype=np.float32)
+    control_points = np.ascontiguousarray(control_points, dtype=np.float32)
+    nonlinear_weights = np.ascontiguousarray(nonlinear_weights, dtype=np.float32)
+    affine_weights = np.ascontiguousarray(affine_weights, dtype=np.float32)
+
+    distances = compute_pairwise_distances(target_points, control_points)
+
+    # Ensure kernel matrix is float32
+    kernel_matrix = np.where(
+        distances > 0,
+        distances * cv2.log(distances + 1e-6),
+        0,
+    ).astype(np.float32)
+
+    # Prepare affine terms
+    num_points = len(target_points)
+    affine_terms = np.empty((num_points, 3), dtype=np.float32)
+    affine_terms[:, 0] = 1
+    affine_terms[:, 1:] = target_points
+
+    # Matrix multiplications with consistent float32 type
+    nonlinear_part = cv2.gemm(kernel_matrix, nonlinear_weights, 1, None, 0)
+    affine_part = cv2.gemm(affine_terms, affine_weights, 1, None, 0)
+
+    return nonlinear_part + affine_part
+
+
+def get_camera_matrix_distortion_maps(
+    image_shape: tuple[int, int],
+    k: float,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate distortion maps using camera matrix model.
+
+    Args:
+        image_shape (tuple[int, int]): Image shape (height, width)
+        k (float): Distortion coefficient
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: Tuple of (map_x, map_y) distortion maps
+
+    """
+    height, width = image_shape[:2]
+
+    center_x, center_y = width / 2, height / 2
+
+    camera_matrix = np.array(
+        [[width, 0, center_x], [0, height, center_y], [0, 0, 1]],
+        dtype=np.float32,
+    )
+    distortion = np.array([k, k, 0, 0, 0], dtype=np.float32)
+    return cv2.initUndistortRectifyMap(
+        camera_matrix,
+        distortion,
+        None,
+        None,
+        (width, height),
+        cv2.CV_32FC1,
+    )
+
+
+def get_fisheye_distortion_maps(
+    image_shape: tuple[int, int],
+    k: float,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate distortion maps using fisheye model.
+
+    Args:
+        image_shape (tuple[int, int]): Image shape (height, width)
+        k (float): Distortion coefficient
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: Tuple of (map_x, map_y) distortion maps
+
+    """
+    height, width = image_shape[:2]
+
+    center_x, center_y = width / 2, height / 2
+    # 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)
+
+    # 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) + center_x
+    map_y = r_dist * np.sin(theta) + center_y
+
+    return map_x, map_y
+
+
+def generate_control_points(num_control_points: int) -> np.ndarray:
+    """Generate control points for TPS transformation.
+
+    Args:
+        num_control_points (int): Number of control points per side
+
+    Returns:
+        np.ndarray: Control points with shape (N, 2)
+
+    """
+    if num_control_points == 2:
+        # Generate 4 corners + center point similar to Kornia
+        return np.array(
+            [
+                [0, 0],  # top-left
+                [0, 1],  # bottom-left
+                [1, 0],  # top-right
+                [1, 1],  # bottom-right
+                [0.5, 0.5],  # center
+            ],
+            dtype=np.float32,
+        )
+
+        # Generate regular grid
+    x = np.linspace(0, 1, num_control_points)
+    y = np.linspace(0, 1, num_control_points)
+    return np.stack(np.meshgrid(x, y), axis=-1).reshape(-1, 2)
+
+
+def volume_hflip(volume: np.ndarray) -> np.ndarray:
+    """Perform horizontal flip on a volume (numpy array).
+
+    Flips the volume along the width axis (axis=2). Handles inputs with
+    shapes (D, H, W) or (D, H, W, C).
+
+    Args:
+        volume (np.ndarray): Input volume.
+
+    Returns:
+        np.ndarray: Horizontally flipped volume.
+
+    """
+    return np.flip(volume, axis=2)
+
+
+def volume_vflip(volume: np.ndarray) -> np.ndarray:
+    """Perform vertical flip on a volume (numpy array).
+
+    Flips the volume along the height axis (axis=1). Handles inputs with
+    shapes (D, H, W) or (D, H, W, C).
+
+    Args:
+        volume (np.ndarray): Input volume.
+
+    Returns:
+        np.ndarray: Vertically flipped volume.
+
+    """
+    return np.flip(volume, axis=1)
+
+
+def volumes_hflip(volumes: np.ndarray) -> np.ndarray:
+    """Perform horizontal flip on a batch of volumes (numpy array).
+
+    Flips the volumes along the width axis (axis=3). Handles inputs with
+    shapes (B, D, H, W) or (B, D, H, W, C).
+
+    Args:
+        volumes (np.ndarray): Input batch of volumes.
+
+    Returns:
+        np.ndarray: Horizontally flipped batch of volumes.
+
+    """
+    # Width axis is 3 for both (B, D, H, W) and (B, D, H, W, C)
+    return np.flip(volumes, axis=3)
+
+
+def volumes_vflip(volumes: np.ndarray) -> np.ndarray:
+    """Perform vertical flip on a batch of volumes (numpy array).
+
+    Flips the volumes along the height axis (axis=2). Handles inputs with
+    shapes (B, D, H, W) or (B, D, H, W, C).
+
+    Args:
+        volumes (np.ndarray): Input batch of volumes.
+
+    Returns:
+        np.ndarray: Vertically flipped batch of volumes.
+
+    """
+    # Height axis is 2 for both (B, D, H, W) and (B, D, H, W, C)
+    return np.flip(volumes, axis=2)
+
+
+def volume_rot90(volume: np.ndarray, factor: Literal[0, 1, 2, 3]) -> np.ndarray:
+    """Rotate a volume 90 degrees counter-clockwise multiple times.
+
+    Rotates the volume in the height-width plane (axes 1 and 2).
+    Handles inputs with shapes (D, H, W) or (D, H, W, C).
+
+    Args:
+        volume (np.ndarray): Input volume.
+        factor (Literal[0, 1, 2, 3]): Number of 90-degree rotations.
+
+    Returns:
+        np.ndarray: Rotated volume.
+
+    """
+    # Axes 1 (height) and 2 (width) for rotation
+    return np.rot90(volume, k=factor, axes=(1, 2))
+
+
+def volumes_rot90(volumes: np.ndarray, factor: Literal[0, 1, 2, 3]) -> np.ndarray:
+    """Rotate a batch of volumes 90 degrees counter-clockwise multiple times.
+
+    Rotates the volumes in the height-width plane (axes 2 and 3).
+    Handles inputs with shapes (B, D, H, W) or (B, D, H, W, C).
+
+    Args:
+        volumes (np.ndarray): Input batch of volumes.
+        factor (Literal[0, 1, 2, 3]): Number of 90-degree rotations
+
+    Returns:
+        np.ndarray: Rotated batch of volumes.
+
+    """
+    # Axes 2 (height) and 3 (width) for rotation
+    return np.rot90(volumes, k=factor, axes=(2, 3))
+
+
+@preserve_channel_dim
+def erode(img: np.ndarray, kernel: np.ndarray) -> np.ndarray:
+    """Apply erosion to an image.
+
+    This function applies erosion to an image using the cv2.erode function.
+
+    Args:
+        img (np.ndarray): Input image as a numpy array.
+        kernel (np.ndarray): Kernel as a numpy array.
+
+    Returns:
+        np.ndarray: The eroded image.
+
+    """
+    return cv2.erode(img, kernel, iterations=1)
+
+
+@preserve_channel_dim
+def dilate(img: np.ndarray, kernel: np.ndarray) -> np.ndarray:
+    """Apply dilation to an image.
+
+    This function applies dilation to an image using the cv2.dilate function.
+
+    Args:
+        img (np.ndarray): Input image as a numpy array.
+        kernel (np.ndarray): Kernel as a numpy array.
+
+    Returns:
+        np.ndarray: The dilated image.
+
+    """
+    return cv2.dilate(img, kernel, iterations=1)
+
+
+def morphology(
+    img: np.ndarray,
+    kernel: np.ndarray,
+    operation: Literal["dilation", "erosion"],
+) -> np.ndarray:
+    """Apply morphology to an image.
+
+    This function applies morphology to an image using the cv2.morphologyEx function.
+
+    Args:
+        img (np.ndarray): Input image as a numpy array.
+        kernel (np.ndarray): Kernel as a numpy array.
+        operation (Literal["dilation", "erosion"]): The operation to apply.
+
+    Returns:
+        np.ndarray: The morphology applied to the image.
+
+    """
+    if operation == "dilation":
+        return dilate(img, kernel)
+    if operation == "erosion":
+        return erode(img, kernel)
+
+    raise ValueError(f"Unsupported operation: {operation}")
+
+
+@handle_empty_array("bboxes")
+def bboxes_morphology(
+    bboxes: np.ndarray,
+    kernel: np.ndarray,
+    operation: Literal["dilation", "erosion"],
+    image_shape: tuple[int, int],
+) -> np.ndarray:
+    """Apply morphology to bounding boxes.
+
+    This function applies morphology to bounding boxes by first converting the bounding
+    boxes to a mask and then applying the morphology to the mask.
+
+    Args:
+        bboxes (np.ndarray): Bounding boxes as a numpy array.
+        kernel (np.ndarray): Kernel as a numpy array.
+        operation (Literal["dilation", "erosion"]): The operation to apply.
+        image_shape (tuple[int, int]): The shape of the image.
+
+    Returns:
+        np.ndarray: The morphology applied to the bounding boxes.
+
+    """
+    bboxes = bboxes.copy()
+    masks = masks_from_bboxes(bboxes, image_shape)
+    masks = morphology(masks, kernel, operation)
+    bboxes[:, :4] = bboxes_from_masks(masks)
+    return bboxes
+
+
+D4_TRANSFORMATIONS_IMAGES = {
+    "e": lambda x: x,  # Identity transformation
+    "r90": lambda x: rot90_images(x, 1),  # Rotate 90 degrees
+    "r180": lambda x: rot90_images(x, 2),  # Rotate 180 degrees
+    "r270": lambda x: rot90_images(x, 3),  # Rotate 270 degrees
+    "v": vflip,  # Vertical flip (already batch-aware)
+    "hvt": lambda x: transpose_images(rot90_images(x, 2)),  # Reflect over anti-diagonal
+    "h": hflip,  # Horizontal flip (already batch-aware)
+    "t": transpose_images,  # Transpose (reflect over main diagonal)
+}
+
+
+def d4_images(img: np.ndarray, group_member: Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"]) -> np.ndarray:
+    """Applies a `D_4` symmetry group transformation to a batch of images.
+
+    This function manipulates a batch of images using transformations such as rotations and flips,
+    corresponding to the `D_4` dihedral group symmetry operations.
+    Each transformation is identified by a unique group member code.
+
+    Args:
+        img (np.ndarray): The input batch of images to transform with shape:
+            - (N, H, W) for grayscale images
+            - (N, H, W, C) for multi-channel images
+            where N is the batch size, H is height, W is width, C is channels
+        group_member (Literal["e", "r90", "r180", "r270", "v", "hvt", "h", "t"]): A string identifier indicating
+            the specific transformation to apply. Valid codes include:
+            - 'e': Identity (no transformation).
+            - 'r90': Rotate 90 degrees counterclockwise.
+            - 'r180': Rotate 180 degrees.
+            - 'r270': Rotate 270 degrees counterclockwise.
+            - 'v': Vertical flip.
+            - 'hvt': Transpose over second diagonal
+            - 'h': Horizontal flip.
+            - 't': Transpose (reflect over the main diagonal).
+
+    Returns:
+        np.ndarray: The transformed batch of images.
+
+    """
+    # Execute the appropriate transformation
+    return D4_TRANSFORMATIONS_IMAGES[group_member](img)