Geometric functional transforms (augmentations.geometric.functional)¶

def adjust_padding_by_position(
    h_top: int,
    h_bottom: int,
    w_left: int,
    w_right: int,
    position: PositionType,
    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 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.

    Example:
        >>> 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)])

@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:]

@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)

@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_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.

    Example:
        >>> 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_d4(
    bboxes: np.ndarray,
    group_member: D4Type,
) -> 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.

    Parameters:
    -  bboxes: 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 (D4Type): A string identifier for the `D_4` group transformation to apply.
        Valid values are 'e', 'r90', 'r180', 'r270', 'v', 'hvt', 'h', 't'.

    Returns:
    - BoxInternalType: The transformed bounding box.

    Raises:
    - ValueError: If an invalid group member is specified.

    Examples:
    - Applying a 90-degree rotation:
      `bbox_d4((10, 20, 110, 120), 'r90')`
      This would rotate the bounding box 90 degrees within a 100x100 image.
    """
    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("bboxes")
def bboxes_flip(bboxes: np.ndarray, d: int) -> np.ndarray:
    """Flip a bounding box either vertically, horizontally or both depending on the value of `d`.

    Args:
        bboxes: 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, ...).
        d: dimension. 0 for vertical flip, 1 for horizontal, -1 for transpose

    Returns:
        A bounding box `(x_min, y_min, x_max, y_max)`.

    Raises:
        ValueError: if value of `d` is not -1, 0 or 1.

    """
    if d == 0:
        return bboxes_vflip(bboxes)
    if d == 1:
        return bboxes_hflip(bboxes)
    if d == -1:
        bboxes = bboxes_hflip(bboxes)
        return bboxes_vflip(bboxes)

    raise ValueError(f"Invalid d value {d}. Valid values are -1, 0 and 1")

@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:
    """Apply grid shuffle transformation to bounding boxes.

    This function transforms bounding boxes according to a grid shuffle operation. It handles cases
    where bounding boxes may be split into multiple components after shuffling and applies
    filtering based on minimum area and visibility requirements.

    Args:
        bboxes: Array of bounding boxes with shape (N, 4+) where N is the number of boxes.
               Each box is in format [x_min, y_min, x_max, y_max, ...], where ... represents
               optional additional fields (e.g., class_id, score).
        tiles: Array of tile coordinates with shape (M, 4) where M is the number of tiles.
               Each tile is in format [start_y, start_x, end_y, end_x].
        mapping: List of indices defining how tiles should be rearranged. Each index i in the list
                contains the index of the tile that should be moved to position i.
        image_shape: Shape of the image as (height, width).
        min_area: Minimum area threshold in pixels. If a component's area after shuffling is
                 smaller than this value, it will be filtered out. If None, no area filtering
                 is applied.
        min_visibility: Minimum visibility ratio threshold in range [0, 1]. Calculated as
                       (component_area / original_area). If a component's visibility is lower
                       than this value, it will be filtered out. If None, no visibility
                       filtering is applied.

    Returns:
        np.ndarray: Array of transformed bounding boxes with shape (K, 4+) where K is the
                   number of valid components after shuffling and filtering. The format of
                   each box matches the input format, preserving any additional fields.
                   If no valid components remain after filtering, returns an empty array
                   with shape (0, C) where C matches the input column count.

    Note:
        - The function converts bboxes to masks before applying the transformation to handle
          cases where boxes may be split into multiple components.
        - After shuffling, each component is validated against min_area and min_visibility
          requirements independently.
        - Additional bbox fields (beyond x_min, y_min, x_max, y_max) are preserved and
          copied to all components derived from the same original bbox.
        - Empty input arrays are handled gracefully and return empty arrays of the
          appropriate shape.

    Example:
        >>> bboxes = np.array([[10, 10, 90, 90]])  # Single box crossing multiple tiles
        >>> tiles = np.array([
        ...     [0, 0, 50, 50],    # top-left tile
        ...     [0, 50, 50, 100],  # top-right tile
        ...     [50, 0, 100, 50],  # bottom-left tile
        ...     [50, 50, 100, 100] # bottom-right tile
        ... ])
        >>> mapping = [3, 2, 1, 0]  # Rotate tiles counter-clockwise
        >>> result = bboxes_grid_shuffle(bboxes, tiles, mapping, (100, 100), 100, 0.2)
        >>> # Result may contain multiple boxes if the original box was split
    """
    # 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

@handle_empty_array("bboxes")
def bboxes_hflip(bboxes: np.ndarray) -> np.ndarray:
    """Flip bounding boxes horizontally around the y-axis.

    Args:
        bboxes: 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, ...).

    Returns:
        np.ndarray: A numpy array of horizontally flipped bounding boxes with the same shape as input.
    """
    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_rot90(bboxes: np.ndarray, factor: int) -> np.ndarray:
    """Rotates bounding boxes by 90 degrees CCW (see np.rot90)

    Args:
        bboxes: 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, ...).
        factor: Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.

    Returns:
        np.ndarray: A numpy array of rotated bounding boxes with the same shape as input.

    Raises:
        ValueError: If factor is not in set {0, 1, 2, 3}.
    """
    if factor not in {0, 1, 2, 3}:
        raise ValueError("Parameter factor must be in set {0, 1, 2, 3}")

    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_transpose(bboxes: np.ndarray) -> np.ndarray:
    """Transpose bounding boxes by swapping x and y coordinates.

    Args:
        bboxes: 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, ...).

    Returns:
        np.ndarray: A numpy array of transposed bounding boxes with the same shape as input.
    """
    transposed_bboxes = bboxes.copy()
    transposed_bboxes[:, [0, 1, 2, 3]] = bboxes[:, [1, 0, 3, 2]]

    return transposed_bboxes

@handle_empty_array("bboxes")
def bboxes_vflip(bboxes: np.ndarray) -> np.ndarray:
    """Flip bounding boxes vertically around the x-axis.

    Args:
        bboxes: 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, ...).

    Returns:
        np.ndarray: A numpy array of vertically flipped bounding boxes with the same shape as input.
    """
    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

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

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 compute_tps_weights(src_points: np.ndarray, dst_points: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Compute Thin Plate Spline weights.

    Args:
        src_points: Source control points with shape (num_points, 2)
        dst_points: Destination control points with shape (num_points, 2)

    Returns:
        tuple of:
        - 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 = np.linalg.norm(src_points[:, None] - src_points, axis=2)

    # Apply TPS kernel function: U(r) = r² log(r)
    # Add small epsilon to avoid log(0)
    kernel_matrix = np.where(distances > 0, distances * distances * np.log(distances + 1e-6), 0)

    # Construct affine terms matrix [1, x, y]
    affine_terms = np.ones((num_points, 3))
    affine_terms[:, 1:] = src_points

    # Build system matrix
    system_matrix = np.zeros((num_points + 3, num_points + 3))
    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

    # Right-hand side of the system
    target_coords = np.zeros((num_points + 3, 2))
    target_coords[:num_points] = dst_points

    # Solve the system for both x and y coordinates
    all_weights = np.linalg.solve(system_matrix, target_coords)

    # Split weights into nonlinear and affine components
    nonlinear_weights = all_weights[:num_points]
    affine_weights = all_weights[num_points:]

    return nonlinear_weights, affine_weights

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 create_affine_transformation_matrix(
    translate: XYInt,
    shear: XYFloat,
    scale: XYFloat,
    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 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."""
    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

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