Helper functions for working with keypoints (augmentations.core.keypoints_utils)¶

def check_keypoints(keypoints: np.ndarray, shape: ShapeType) -> None:
    """Check if keypoint coordinates are within valid ranges for the given shape.

    This function validates that:
    1. All x-coordinates are within [0, width)
    2. All y-coordinates are within [0, height)
    3. If depth is provided in shape, z-coordinates are within [0, depth)
    4. Angles are within the range [0, 2π)

    Args:
        keypoints (np.ndarray): Array of keypoints with shape (N, 5+), where N is the number of keypoints.
            - First 2 columns are always x, y
            - Column 3 (if present) is z
            - Column 4 (if present) is angle
            - Column 5+ (if present) are additional attributes
        shape (ShapeType): The shape of the image/volume:
                           - For 2D: {'height': int, 'width': int}
                           - For 3D: {'height': int, 'width': int, 'depth': int}

    Raises:
        ValueError: If any keypoint coordinate is outside the valid range, or if angles are invalid.
                   The error message will detail which keypoints are invalid and why.

    Note:
        - The function assumes that keypoint coordinates are in absolute pixel values, not normalized
        - Angles are in radians
        - Z-coordinates are only checked if 'depth' is present in shape
    """
    height, width = shape["height"], shape["width"]
    has_depth = "depth" in shape

    # Check x and y coordinates (always present)
    x, y = keypoints[:, 0], keypoints[:, 1]
    invalid_x = np.where((x < 0) | (x >= width))[0]
    invalid_y = np.where((y < 0) | (y >= height))[0]

    error_messages = []

    # Handle x, y errors
    for idx in sorted(set(invalid_x) | set(invalid_y)):
        if idx in invalid_x:
            error_messages.append(
                f"Expected x for keypoint {keypoints[idx]} to be in range [0, {width}), got {x[idx]}",
            )
        if idx in invalid_y:
            error_messages.append(
                f"Expected y for keypoint {keypoints[idx]} to be in range [0, {height}), got {y[idx]}",
            )

    # Check z coordinates if depth is provided and keypoints have z
    if has_depth and keypoints.shape[1] > 2:
        z = keypoints[:, 2]
        depth = shape["depth"]
        invalid_z = np.where((z < 0) | (z >= depth))[0]
        error_messages.extend(
            f"Expected z for keypoint {keypoints[idx]} to be in range [0, {depth}), got {z[idx]}" for idx in invalid_z
        )

    # Check angles only if keypoints have angle column
    if keypoints.shape[1] > 3:
        angles = keypoints[:, 3]
        invalid_angles = np.where((angles < 0) | (angles >= 2 * math.pi))[0]
        error_messages.extend(
            f"Expected angle for keypoint {keypoints[idx]} to be in range [0, 2π), got {angles[idx]}"
            for idx in invalid_angles
        )

    if error_messages:
        raise ValueError("\n".join(error_messages))

def convert_keypoints_from_albumentations(
    keypoints: np.ndarray,
    target_format: Literal["xy", "yx", "xya", "xys", "xyas", "xysa", "xyz"],
    shape: ShapeType,
    check_validity: bool = False,
    angle_in_degrees: bool = True,
) -> np.ndarray:
    """Convert keypoints from Albumentations format to various other formats.

    This function takes keypoints in the standard Albumentations format [x, y, z, angle, scale]
    and converts them to the specified target format.

    Args:
        keypoints (np.ndarray): Array of keypoints in Albumentations format with shape (N, 5+),
                                where N is the number of keypoints. Each row represents a keypoint
                                [x, y, z, angle, scale, ...].
        target_format (Literal["xy", "yx", "xya", "xys", "xyas", "xysa", "xyz"]): The desired output format.
            - "xy": [x, y]
            - "yx": [y, x]
            - "xya": [x, y, angle]
            - "xys": [x, y, scale]
            - "xyas": [x, y, angle, scale]
            - "xysa": [x, y, scale, angle]
            - "xyz": [x, y, z]
        shape (ShapeType): The shape of the image {'height': height, 'width': width, 'depth': depth}.
        check_validity (bool, optional): If True, check if the keypoints are within the image boundaries.
                                         Defaults to False.
        angle_in_degrees (bool, optional): If True, convert output angles to degrees.
                                           If False, angles remain in radians.
                                           Defaults to True.

    Returns:
        np.ndarray: Array of keypoints in the specified target format with shape (N, 2+).
                    Any additional columns from the input keypoints beyond the first 5
                    are preserved and appended after the converted columns.

    Raises:
        ValueError: If the target_format is not one of the supported formats.

    Note:
        - Input angles are assumed to be in the range [0, 2π) radians
        - If the input keypoints have additional columns beyond the first 5,
          these columns are preserved in the output
    """
    if target_format not in keypoint_formats:
        raise ValueError(f"Unknown target_format {target_format}. Supported formats are: {keypoint_formats}")

    x, y, z, angle, scale = keypoints[:, 0], keypoints[:, 1], keypoints[:, 2], keypoints[:, 3], keypoints[:, 4]
    angle = angle_to_2pi_range(angle)

    if check_validity:
        check_keypoints(np.column_stack((x, y, z, angle, scale)), shape)

    if angle_in_degrees:
        angle = np.degrees(angle)

    format_to_columns = {
        "xy": [x, y],
        "yx": [y, x],
        "xya": [x, y, angle],
        "xys": [x, y, scale],
        "xyas": [x, y, angle, scale],
        "xysa": [x, y, scale, angle],
        "xyz": [x, y, z],
    }

    result = np.column_stack(format_to_columns[target_format])

    # Add any additional columns from the original keypoints
    if keypoints.shape[1] > NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS:
        return np.column_stack((result, keypoints[:, NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS:]))

    return result

def convert_keypoints_to_albumentations(
    keypoints: np.ndarray,
    source_format: Literal["xy", "yx", "xya", "xys", "xyas", "xysa", "xyz"],
    shape: ShapeType,
    check_validity: bool = False,
    angle_in_degrees: bool = True,
) -> np.ndarray:
    """Convert keypoints from various formats to the Albumentations format.

    This function takes keypoints in different formats and converts them to the standard
    Albumentations format: [x, y, z, angle, scale]. For 2D formats, z is set to 0.
    For formats without angle or scale, these values are set to 0.

    Args:
        keypoints (np.ndarray): Array of keypoints with shape (N, 2+), where N is the number of keypoints.
                                The number of columns depends on the source_format.
        source_format (Literal["xy", "yx", "xya", "xys", "xyas", "xysa", "xyz"]): The format of the input keypoints.
            - "xy": [x, y]
            - "yx": [y, x]
            - "xya": [x, y, angle]
            - "xys": [x, y, scale]
            - "xyas": [x, y, angle, scale]
            - "xysa": [x, y, scale, angle]
            - "xyz": [x, y, z]
        shape (ShapeType): The shape of the image {'height': height, 'width': width, 'depth': depth}.
        check_validity (bool, optional): If True, check if the converted keypoints are within the image boundaries.
                                         Defaults to False.
        angle_in_degrees (bool, optional): If True, convert input angles from degrees to radians.
                                           Defaults to True.

    Returns:
        np.ndarray: Array of keypoints in Albumentations format [x, y, z, angle, scale] with shape (N, 5+).
                    Any additional columns from the input keypoints are preserved and appended after the
                    first 5 columns.

    Raises:
        ValueError: If the source_format is not one of the supported formats.

    Note:
        - For 2D formats (xy, yx, xya, xys, xyas, xysa), z coordinate is set to 0
        - Angles are converted to the range [0, 2π) radians
        - If the input keypoints have additional columns beyond what's specified in the source_format,
          these columns are preserved in the output
    """
    if source_format not in keypoint_formats:
        raise ValueError(f"Unknown source_format {source_format}. Supported formats are: {keypoint_formats}")

    format_to_indices: dict[str, list[int | None]] = {
        "xy": [0, 1, None, None, None],
        "yx": [1, 0, None, None, None],
        "xya": [0, 1, None, 2, None],
        "xys": [0, 1, None, None, 2],
        "xyas": [0, 1, None, 2, 3],
        "xysa": [0, 1, None, 3, 2],
        "xyz": [0, 1, 2, None, None],
    }

    indices: list[int | None] = format_to_indices[source_format]

    processed_keypoints = np.zeros((keypoints.shape[0], NUM_KEYPOINTS_COLUMNS_IN_ALBUMENTATIONS), dtype=np.float32)

    for i, idx in enumerate(indices):
        if idx is not None:
            processed_keypoints[:, i] = keypoints[:, idx]

    if angle_in_degrees and indices[3] is not None:  # angle is now at index 3
        processed_keypoints[:, 3] = np.radians(processed_keypoints[:, 3])

    processed_keypoints[:, 3] = angle_to_2pi_range(processed_keypoints[:, 3])  # angle is now at index 3

    if keypoints.shape[1] > len(source_format):
        processed_keypoints = np.column_stack((processed_keypoints, keypoints[:, len(source_format) :]))

    if check_validity:
        check_keypoints(processed_keypoints, shape)

    return processed_keypoints

def filter_keypoints(
    keypoints: np.ndarray,
    shape: ShapeType,
    remove_invisible: bool,
) -> np.ndarray:
    """Filter keypoints to remove those outside the boundaries.

    Args:
        keypoints: A numpy array of shape (N, 5+) where N is the number of keypoints.
                   Each row represents a keypoint (x, y, z, angle, scale, ...).
        shape: Shape to check against as {'height': height, 'width': width, 'depth': depth}.
        remove_invisible: If True, remove keypoints outside the boundaries.

    Returns:
        A numpy array of filtered keypoints.
    """
    if not remove_invisible:
        return keypoints

    if not keypoints.size:
        return keypoints

    height, width, depth = shape["height"], shape["width"], shape.get("depth", None)

    # Create boolean mask for visible keypoints
    x, y, z = keypoints[:, 0], keypoints[:, 1], keypoints[:, 2]
    visible = (x >= 0) & (x < width) & (y >= 0) & (y < height)

    if depth is not None:
        visible &= (z >= 0) & (z < depth)

    # Apply the mask to filter keypoints
    return keypoints[visible]