Function make_nifti_image

tests/test_utils.py:387–418

Create a temporary nifti image on the disk and return the image name. User is responsible for deleting the temporary file when done with it.

(
    array: NdarrayOrTensor, affine=None, dir=None, fname=None, suffix=".nii.gz", verbose=False, dtype=float
)

def make_nifti_image(
    array: NdarrayOrTensor, affine=None, dir=None, fname=None, suffix=".nii.gz", verbose=False, dtype=float
):
    """
    Create a temporary nifti image on the disk and return the image name.
    User is responsible for deleting the temporary file when done with it.
    """
    if isinstance(array, torch.Tensor):
        array, *_ = convert_data_type(array, np.ndarray)
    if isinstance(affine, torch.Tensor):
        affine, *_ = convert_data_type(affine, np.ndarray)
    if affine is None:
        affine = np.eye(4)
    test_image = nib.Nifti1Image(array.astype(dtype), affine)  # type: ignore

    # if dir not given, create random. Else, make sure it exists.
    if dir is None:
        dir = tempfile.mkdtemp()
    else:
        os.makedirs(dir, exist_ok=True)

    # If fname not given, get random one. Else, concat dir, fname and suffix.
    if fname is None:
        temp_f, fname = tempfile.mkstemp(suffix=suffix, dir=dir)
        os.close(temp_f)
    else:
        fname = os.path.join(dir, fname + suffix)

    nib.save(test_image, fname)
    if verbose:
        print(f"File written: {fname}.")
    return fname


def make_rand_affine(ndim: int = 3, random_state: np.random.RandomState | None = None):