Function map_spatial_axes

monai/transforms/utils.py:1684–1716

Utility to map the spatial axes to real axes in channel first/last shape. For example: If `channel_first` is True, and `img` has 3 spatial dims, map spatial axes to real axes as below: None -> [1, 2, 3] [0, 1] -> [1, 2] [0, -1] -> [1, -1] If `channel_first` is False, and

(
    img_ndim: int, spatial_axes: Sequence[int] | int | None = None, channel_first: bool = True
)

def map_spatial_axes(
    img_ndim: int, spatial_axes: Sequence[int] | int | None = None, channel_first: bool = True
) -> list[int]:
    """
    Utility to map the spatial axes to real axes in channel first/last shape.
    For example:
    If `channel_first` is True, and `img` has 3 spatial dims, map spatial axes to real axes as below:
    None -> [1, 2, 3]
    [0, 1] -> [1, 2]
    [0, -1] -> [1, -1]
    If `channel_first` is False, and `img` has 3 spatial dims, map spatial axes to real axes as below:
    None -> [0, 1, 2]
    [0, 1] -> [0, 1]
    [0, -1] -> [0, -2]

    Args:
        img_ndim: dimension number of the target image.
        spatial_axes: spatial axes to be converted, default is None.
            The default `None` will convert to all the spatial axes of the image.
            If axis is negative it counts from the last to the first axis.
            If axis is a tuple of ints.
        channel_first: the image data is channel first or channel last, default to channel first.

    """
    if spatial_axes is None:
        return list(range(1, img_ndim) if channel_first else range(img_ndim - 1))
    spatial_axes_ = []
    for a in ensure_tuple(spatial_axes):
        if channel_first:
            spatial_axes_.append(a % img_ndim if a < 0 else a + 1)
        else:
            spatial_axes_.append((a - 1) % (img_ndim - 1) if a < 0 else a)
    return spatial_axes_


@contextmanager