Class NibabelReader

monai/data/image_reader.py:1001–1221

Load NIfTI format images based on Nibabel library. Args: channel_dim: the channel dimension of the input image, default is None. this is used to set original_channel_dim in the metadata, EnsureChannelFirstD reads this field. if None, `original_channel_dim` w

@require_pkg(pkg_name="nibabel")
class NibabelReader(ImageReader):
    """
    Load NIfTI format images based on Nibabel library.

    Args:
        channel_dim: the channel dimension of the input image, default is None.
            this is used to set original_channel_dim in the metadata, EnsureChannelFirstD reads this field.
            if None, `original_channel_dim` will be either `no_channel` or `-1`.
            most Nifti files are usually "channel last", no need to specify this argument for them.
        as_closest_canonical: if True, load the image as closest to canonical axis format.
        squeeze_non_spatial_dims: if True, non-spatial singletons will be squeezed, e.g. (256,256,1,3) -> (256,256,3)
        to_gpu: If True, load the image into GPU memory using CuPy and Kvikio. This can accelerate data loading.
            Default is False. CuPy and Kvikio are required for this option.
            Note: For compressed NIfTI files, some operations may still be performed on CPU memory,
            and the acceleration may not be significant. In some cases, it may be slower than loading on CPU.
        kwargs: additional args for `nibabel.load` API. more details about available args:
            https://github.com/nipy/nibabel/blob/master/nibabel/loadsave.py

    """

    def __init__(
        self,
        channel_dim: str | int | None = None,
        as_closest_canonical: bool = False,
        squeeze_non_spatial_dims: bool = False,
        to_gpu: bool = False,
        **kwargs,
    ):
        super().__init__()
        self.channel_dim = float("nan") if channel_dim == "no_channel" else channel_dim
        self.as_closest_canonical = as_closest_canonical
        self.squeeze_non_spatial_dims = squeeze_non_spatial_dims
        if to_gpu and (not has_cp or not has_kvikio):
            warnings.warn(
                "NibabelReader: CuPy and/or Kvikio not installed for GPU loading, falling back to CPU loading."
            )
            to_gpu = False

        if to_gpu:
            self.warmup_kvikio()

        self.to_gpu = to_gpu
        self.kwargs = kwargs

    def warmup_kvikio(self):
        """
        Warm up the Kvikio library to initialize the internal buffers, cuFile, GDS, etc.
        This can accelerate the data loading process when `to_gpu` is set to True.
        """
        if has_cp and has_kvikio:
            a = cp.arange(100)
            with tempfile.NamedTemporaryFile() as tmp_file:
                tmp_file_name = tmp_file.name
                f = kvikio.CuFile(tmp_file_name, "w")
                f.write(a)
                f.close()

                b = cp.empty_like(a)