nifti.py 文件源码

def write_nifti(data, output_fname, header=None, affine=None, use_data_dtype=True, **kwargs):
    """Write data to a nifti file.

    This will write the output directory if it does not exist yet.

    Args:
        data (ndarray): the data to write to that nifti file
        output_fname (str): the name of the resulting nifti file, this function will append .nii.gz if no
            suitable extension is given.
        header (nibabel header): the nibabel header to use as header for the nifti file. If None we will use
            a default header.
        affine (ndarray): the affine transformation matrix
        use_data_dtype (boolean): if we want to use the dtype from the data instead of that from the header
            when saving the nifti.
        **kwargs: other arguments to Nifti2Image from NiBabel
    """
    if header is None:
        header = nib.nifti2.Nifti2Header()

    if use_data_dtype:
        header = copy.deepcopy(header)
        dtype = data.dtype
        if data.dtype == np.bool:
            dtype = np.char
        try:
            header.set_data_dtype(dtype)
        except nib.spatialimages.HeaderDataError:
            pass

    if not (output_fname.endswith('.nii.gz') or output_fname.endswith('.nii')):
        output_fname += '.nii.gz'

    if not os.path.exists(os.path.dirname(output_fname)):
        os.makedirs(os.path.dirname(output_fname))

    if isinstance(header, nib.nifti2.Nifti2Header):
        format = nib.Nifti2Image
    else:
        format = nib.Nifti1Image

    format(data, affine, header=header, **kwargs).to_filename(output_fname)