Python nibabel.save() Examples

def generateImgSlicesFolder(data_folder = '../Data/MS2017a/scans/'):
	scan_folders = glob.glob(data_folder + '*')

	for sf in scan_folders:
		slice_dir_path = os.path.join(sf, 'slices/')
		if not os.path.exists(slice_dir_path):
			print('Creating directory at:' , slice_dir_path)
			os.makedirs(slice_dir_path)

		img = nib.load(os.path.join(sf, 'pre/FLAIR.nii.gz'))
		img_np = img.get_data()
		img_affine = img.affine
		print(sf)
		print('The img shape', img_np.shape[2])
		for i in range(img_np.shape[2]):
			slice_img_np = img_np[:,:,i]
			nft_img = nib.Nifti1Image(slice_img_np, img_affine)
			nib.save(nft_img, slice_dir_path + 'FLAIR_' + str(i) + '.nii.gz')

			if os.path.basename(sf) == '0':
				slice_img = nib.load(slice_dir_path + 'FLAIR_' + str(i) + '.nii.gz').get_data() / 5
				print('DID I GET HERE?')
				print('Writing to', str(i) + '.jpg') 

def generateGTSlicesFolder(data_folder = '../Data/MS2017a/scans/'):
	scan_folders = glob.glob(data_folder + '*')

	for sf in scan_folders:
		slice_dir_path = os.path.join(sf, 'gt_slices/')
		if not os.path.exists(slice_dir_path):
			print('Creating directory at:' , slice_dir_path)
			os.makedirs(slice_dir_path)

		img = nib.load(os.path.join(sf, 'wmh.nii.gz'))
		img_np = img.get_data()
		img_affine = img.affine
		print(sf)
		print('The img shape', img_np.shape[2])
		for i in range(img_np.shape[2]):
			slice_img_np = img_np[:,:,i]
			nft_img = nib.Nifti1Image(slice_img_np, img_affine)
			nib.save(nft_img, slice_dir_path + 'wmh_' + str(i) + '.nii.gz')

			if os.path.basename(sf) == '0':
				slice_img = nib.load(slice_dir_path + 'wmh_' + str(i) + '.nii.gz').get_data() * 256
				#cv2.imwrite('temp/' + str(i) + '.jpg', slice_img) 

def _write_gifti(pd, opth):
    # TODO: what about pointdata?
    from nibabel.gifti.gifti import GiftiDataArray

    if not pd.has_only_triangle:
        raise ValueError('GIFTI writer only accepts triangles.')

    points = GiftiDataArray(data=pd.Points, intent=INTENT_POINTS)
    cells = GiftiDataArray(data=pd.GetCells2D(), intent=INTENT_CELLS)
    # if data is not None:
    #     data_array = GiftiDataArray(data=data, intent=INTENT_POINTDATA)
    #     gii = nb.gifti.GiftiImage(darrays=[points, cells, data_array])
    # else:
    g = nb.gifti.GiftiImage(darrays=[points, cells])
    nb.save(g, opth)


###############################################################################
# VTK Reader and Writer for GIFTI surfaces
############################################################################### 

def test_shape(self, num_workers, dataset_size):
        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
        tempdir = tempfile.mkdtemp()
        nib.save(test_image, os.path.join(tempdir, "test_image1.nii.gz"))
        nib.save(test_image, os.path.join(tempdir, "test_label1.nii.gz"))
        nib.save(test_image, os.path.join(tempdir, "test_extra1.nii.gz"))
        test_data = [
            {
                "image": os.path.join(tempdir, "test_image1.nii.gz"),
                "label": os.path.join(tempdir, "test_label1.nii.gz"),
                "extra": os.path.join(tempdir, "test_extra1.nii.gz"),
            }
        ] * dataset_size
        dataset = CacheDataset(
            data=test_data,
            transform=Compose([LoadNiftid(keys=["image", "label", "extra"])]),
            cache_rate=1,
            num_workers=num_workers,
        )
        shutil.rmtree(tempdir)
        self.assertEqual(len(dataset._cache), dataset.cache_num)
        for i in range(dataset.cache_num):
            self.assertIsNotNone(dataset._cache[i]) 

def test_register_dwi():
    fdata, fbval, fbvec = dpd.get_fnames('small_64D')
    with nbtmp.InTemporaryDirectory() as tmpdir:
        # Use an abbreviated data-set:
        img = nib.load(fdata)
        data = img.get_fdata()[..., :10]
        nib.save(nib.Nifti1Image(data, img.affine),
                 op.join(tmpdir, 'data.nii.gz'))
        # Save a subset:
        bvals = np.loadtxt(fbval)
        bvecs = np.loadtxt(fbvec)
        np.savetxt(op.join(tmpdir, 'bvals.txt'), bvals[:10])
        np.savetxt(op.join(tmpdir, 'bvecs.txt'), bvecs[:10])
        reg_file = register_dwi(op.join(tmpdir, 'data.nii.gz'),
                                op.join(tmpdir, 'bvals.txt'),
                                op.join(tmpdir, 'bvecs.txt'))
        npt.assert_(op.exists(reg_file)) 

def load_or_create(path: PathLike, create: Callable, *args,
                   save: Callable = save, load: Callable = load, **kwargs):
    """
    ``load`` a file from ``path`` if it exists.
    Otherwise ``create`` the value, ``save`` it to ``path``, and return it.

    ``args`` and ``kwargs`` are passed to ``create`` as additional arguments.
    """
    try:
        return load(path)
    except FileNotFoundError:
        pass

    value = create(*args, **kwargs)
    save(value, path)
    return value 

def save_sample(batch_pred, batch_x, batch_y, epoch, batch_id):
    def get_mask(seg_volume):
        seg_volume = seg_volume.cpu().numpy()
        seg_volume = np.squeeze(seg_volume)
        wt_pred = seg_volume[0]
        tc_pred = seg_volume[1]
        et_pred = seg_volume[2]
        mask = np.zeros_like(wt_pred)
        mask[wt_pred > 0.5] = 2
        mask[tc_pred > 0.5] = 1
        mask[et_pred > 0.5] = 4
        mask = mask.astype("uint8")
        mask_nii = nib.Nifti1Image(mask, np.eye(4))
        return mask_nii

    volume = batch_x[:, 0].cpu().numpy()
    volume = volume.astype("uint8")
    volume_nii = nib.Nifti1Image(volume, np.eye(4))
    log_dir = os.path.join(cfg.LOG_DIR, cfg.TASK_NAME, 'epoch'+str(epoch))
    mkdir(log_dir)
    nib.save(volume_nii, os.path.join(log_dir, 'batch'+str(batch_id)+'_volume.nii'))
    pred_nii = get_mask(batch_pred)
    gt_nii = get_mask(batch_y)
    nib.save(pred_nii, os.path.join(log_dir, 'batch' + str(batch_id) + '_pred.nii'))
    nib.save(gt_nii, os.path.join(log_dir, 'batch' + str(batch_id) + '_gt.nii')) 

def generateTrainValFile_Slices(train_fraction, main_folder = '../Data/MS2017a/'):
	train_folders, val_folders = splitTrainVal_Slices(0.8)

	train_folder_names = [train_folders[i].split(main_folder)[1] for i in range(len(train_folders))]
	val_folder_names = [val_folders[i].split(main_folder)[1] for i in range(len(val_folders))]

	f_train = open(main_folder + 'train_slices.txt', 'w+')
	f_val = open(main_folder + 'val_slices.txt', 'w+')

	for fn in train_folder_names:
		f_train.write(fn + '\n')

	for fn in val_folder_names:
		f_val.write(fn + '\n')

	f_train.close()
	f_val.close()

#Use this to save the images quickly (for testing purposes) 

def test_to_nonlinear_grid_crop(self, sphere3_msh):
        import nibabel
        data = sphere3_msh.elm.tag1
        f = mesh_io.ElementData(data, mesh=sphere3_msh)
        affine = np.array([[1, 0, 0, -100.5],
                           [0, 1, 0, -5],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]], dtype=float)
        x, y, z = np.meshgrid(np.arange(-100, 100),
                              np.arange(-5, 5),
                              np.arange(0, 1),
                              indexing='ij')
        nonl_transform = np.concatenate(
            (x[..., None], y[..., None], z[..., None]), axis=3).astype(float)
        img = nibabel.Nifti1Pair(nonl_transform, affine)
        tempf = tempfile.NamedTemporaryFile(suffix='.nii', delete=False)
        fn = tempf.name
        tempf.close()
        nibabel.save(img, fn)
        interp = f.to_deformed_grid(fn, fn, tags=3, method='assign')
        interp = interp.get_data()
        assert np.isclose(interp[100, 5, 0], 3)
        assert np.isclose(interp[187, 5, 0], 0)
        assert np.isclose(interp[193, 5, 0], 0)
        assert np.isclose(interp[198, 5, 0], 0) 

def _reg_prealign(self, row):
        prealign_file = self._get_fname(
            row, '_prealign_from-DWI_to-MNI_xfm.npy')
        if self.force_recompute or not op.exists(prealign_file):
            moving = nib.load(self._b0(row))
            static = afd.read_mni_template()
            moving_data = moving.get_fdata()
            moving_affine = moving.affine
            static_data = static.get_fdata()
            static_affine = static.affine
            _, aff = reg.affine_registration(moving_data,
                                             static_data,
                                             moving_affine,
                                             static_affine)
            np.save(prealign_file, aff)
            meta_fname = self._get_fname(
                row, '_prealign_from-DWI_to-MNI_xfm.json')
            meta = dict(type="rigid")
            afd.write_json(meta_fname, meta)
        return prealign_file 

def _dki(self, row):
        dki_params_file = self._get_fname(row, '_model-DKI_diffmodel.nii.gz')
        if self.force_recompute or not op.exists(dki_params_file):
            img = nib.load(row['dwi_file'])
            data = img.get_fdata()
            gtab = row['gtab']
            brain_mask_file = self._brain_mask(row)
            mask = nib.load(brain_mask_file).get_fdata()
            dkf = dki_fit(gtab, data, mask=mask)
            nib.save(nib.Nifti1Image(dkf.model_params, row['dwi_affine']),
                     dki_params_file)
            meta_fname = self._get_fname(row, '_model-DKI_diffmodel.json')
            meta = dict(
                Parameters=dict(
                    FitMethod="WLS"),
                OutlierRejection=False,
                ModelURL=f"{DIPY_GH}reconst/dki.py")
            afd.write_json(meta_fname, meta)
        return dki_params_file 

def auto_crop_image(input_name, output_name, reserve):
    nim = nib.load(input_name)
    image = nim.get_data()
    X, Y, Z = image.shape[:3]

    # Detect the bounding box of the foreground
    idx = np.nonzero(image > 0)
    x1, x2 = idx[0].min() - reserve, idx[0].max() + reserve + 1
    y1, y2 = idx[1].min() - reserve, idx[1].max() + reserve + 1
    z1, z2 = idx[2].min() - reserve, idx[2].max() + reserve + 1
    x1, x2 = max(x1, 0), min(x2, X)
    y1, y2 = max(y1, 0), min(y2, Y)
    z1, z2 = max(z1, 0), min(z2, Z)
    print('Bounding box')
    print('  bottom-left corner = ({},{},{})'.format(x1, y1, z1))
    print('  top-right corner = ({},{},{})'.format(x2, y2, z2))

    # Crop the image
    image = image[x1:x2, y1:y2, z1:z2]

    # Update the affine matrix
    affine = nim.affine
    affine[:3, 3] = np.dot(affine, np.array([x1, y1, z1, 1]))[:3]
    nim2 = nib.Nifti1Image(image, affine)
    nib.save(nim2, output_name) 

def save_slice_volume(volume, save_path):
    '''
    the function save volume data to slices in the specific directory
    :param volume: input volume data
    :param save_path:
    :return:
    '''
    shape = volume.shape
    # translate intensity to 0-255
    v_max = np.max(volume)
    v_min = np.min(volume)
    volume_norm = (volume - v_min) / (v_max - v_min)
    volume_norm = (volume_norm * 255).astype("int")
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    for i in range(shape[-1]):
        abs_path = os.path.join(save_path, str(i)+".png")
        cv.imwrite(abs_path, volume_norm[..., i])


# calculate the cube information 

def save_image(path_image, image):
    """ save the image into given path

    :param str path_image: path to the image
    :param image: np.array<height, width, ch>

    >>> # Wrong path
    >>> save_image('./missing-path/any-image.png', np.zeros((10, 20)))
    False
    """
    path_dir = os.path.dirname(path_image)
    if not os.path.isdir(path_dir):
        logging.error('upper folder does not exists: "%s"', path_dir)
        return False
    image = convert_ndarray2image(image)
    image.save(path_image) 

def save_landmarks_csv(path_file, landmarks):
    """ save landmarks into a csv file

    we are using simple format::

        ,X,Y
        0,point1-x,point1-y
        1,point2-x,point2-y

    :param str path_file: path to the output file
    :param landmarks: np.array<np_points, dim>
    :return str: file path
    """
    assert os.path.isdir(os.path.dirname(path_file)), \
        'missing folder "%s"' % os.path.dirname(path_file)
    path_file = os.path.splitext(path_file)[0] + '.csv'
    df = pd.DataFrame(landmarks, columns=LANDMARK_COORDS)
    df.index = np.arange(1, len(df) + 1)
    df.to_csv(path_file)
    return path_file 

def make_dti_data(out_fbval, out_fbvec, out_fdata, out_shape=(5, 6, 7)):
    """
    Create a synthetic data-set with a single shell acquisition

    out_fbval, out_fbvec, out_fdata : str
        Full paths to generated data and bval/bvec files

    out_shape : tuple
        The 3D shape of the output volum

    """
    fimg, fbvals, fbvecs = dpd.get_fnames('small_64D')
    img = nib.load(fimg)
    bvals, bvecs = dio.read_bvals_bvecs(fbvals, fbvecs)
    gtab = dpg.gradient_table(bvals, bvecs)

    # Simulate a signal based on the DTI model:
    signal = single_tensor(gtab, S0=100)
    DWI = np.zeros(out_shape + (len(gtab.bvals), ))
    DWI[:] = signal
    nib.save(nib.Nifti1Image(DWI, img.affine), out_fdata)
    np.savetxt(out_fbval, bvals)
    np.savetxt(out_fbvec, bvecs) 

def save_landmarks_pts(path_file, landmarks):
    """ save landmarks into a txt file

    we are using VTK pointdata legacy format, ITK compatible::

        <index, point>
        <number of points>
        point1-x point1-y [point1-z]
        point2-x point2-y [point2-z]

    .. seealso:: https://simpleelastix.readthedocs.io/PointBasedRegistration.html

    :param str path_file: path to the output file
    :param landmarks: np.array<np_points, dim>
    :return str: file path
    """
    assert os.path.isdir(os.path.dirname(path_file)), \
        'missing folder "%s"' % os.path.dirname(path_file)
    path_file = os.path.splitext(path_file)[0] + '.pts'
    lines = ['point', str(len(landmarks))]
    lines += [' '.join(str(i) for i in point) for point in landmarks]
    with open(path_file, 'w') as fp:
        fp.write('\n'.join(lines))
    return path_file 

def test_shape(self, input_param, filenames, expected_shape):
        test_image = np.random.randint(0, 2, size=[128, 128, 128])
        tempdir = tempfile.mkdtemp()
        for i, name in enumerate(filenames):
            filenames[i] = os.path.join(tempdir, name)
            nib.save(nib.Nifti1Image(test_image, np.eye(4)), filenames[i])
        result = LoadNifti(**input_param)(filenames)

        if isinstance(result, tuple):
            result, header = result
            self.assertTrue("affine" in header)
            np.testing.assert_allclose(header["affine"], np.eye(4))
            if input_param["as_closest_canonical"]:
                np.testing.asesrt_allclose(header["original_affine"], np.eye(4))
        self.assertTupleEqual(result.shape, expected_shape)
        shutil.rmtree(tempdir) 

def test_default_none(self, img_transform, expected_shape):
        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4))
        tempdir = tempfile.mkdtemp()
        test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
        test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
        nib.save(test_image, test_image1)
        nib.save(test_image, test_image2)
        test_images = [test_image1, test_image2]
        dataset = ArrayDataset(test_images, img_transform)
        self.assertEqual(len(dataset), 2)
        dataset.set_random_state(1234)
        data1 = dataset[0]
        data2 = dataset[1]
        self.assertTupleEqual(data1.shape, expected_shape)
        self.assertTupleEqual(data2.shape, expected_shape)

        dataset = ArrayDataset(test_images, img_transform)
        dataset.set_random_state(1234)
        _ = dataset[0]
        data2_new = dataset[1]
        np.testing.assert_allclose(data2, data2_new, atol=1e-3)
        shutil.rmtree(tempdir) 

def _write_raw_fif(raw, bids_fname):
    """Save out the raw file in FIF.

    Parameters
    ----------
    raw : mne.io.Raw
        Raw file to save out.
    bids_fname : str | BIDSPath
        The name of the BIDS-specified file where the raw object
        should be saved.

    """
    n_rawfiles = len(raw.filenames)
    if n_rawfiles > 1:
        split_naming = 'bids'
        raw.save(bids_fname, split_naming=split_naming, overwrite=True)
    else:
        # This ensures that single FIF files do not have the split param
        raw.save(bids_fname, split_naming='neuromag', overwrite=True) 

def extractMeanDataStats(size = [200, 200, 100], 
						postfix = '_200x200x100orig', 
						main_folder_path = '../../Data/MS2017b/', 
						):
	scan_folders = glob.glob(main_folder_path + 'scans/*')
	img_path = 'pre/FLAIR' + postfix + '.nii.gz'
	segm_path = 'wmh' + postfix + '.nii.gz'
	
	shape_ = [len(scan_folders), size[0], size[1], size[2]]
	arr = np.zeros(shape_)

	for i, sf in enumerate(scan_folders):
		arr[i, :,:,:] =  numpyFromScan(os.path.join(sf,img_path)).squeeze()

	arr /= len(scan_folders)

	means = np.mean(arr)
	stds = np.std(arr, axis = 0)

	np.save(main_folder_path + 'extra_data/std' + postfix, stds)
	np.save(main_folder_path + 'extra_data/mean' + postfix, means) 

def prepare_example_slice(self, subject_ID):
        "Prepares a coronal slice for the dmipy example notebooks."
        msg = "Preparing coronal slice for dmipy examples"
        print(msg)

        folder_name = "hcp_example_slice"
        example_directory = os.path.join(self.hcp_directory, folder_name)
        if not os.path.exists(example_directory):
            os.makedirs(example_directory)

        subject_data_path = os.path.join(self.hcp_directory, str(subject_ID))

        data = nib.load(os.path.join(
            subject_data_path, 'data.nii.gz')).get_data()
        affine = nib.load(os.path.join(
            subject_data_path, 'data.nii.gz')).affine
        mask = nib.load(os.path.join(
            subject_data_path, 'nodif_brain_mask.nii.gz')).get_data()
        data_shape = data.shape

        slice_index = data_shape[1] // 2
        data_slice = data[:, slice_index: slice_index + 1]
        mask_slice = mask[:, slice_index: slice_index + 1]
        data_slice[mask_slice == 0] = 0

        nib.save(nib.Nifti1Image(data_slice, affine), os.path.join(
            example_directory, 'coronal_slice.nii.gz'))
        print('Done') 

def setUp(self):
        set_determinism(seed=0)

        self.data_dir = tempfile.mkdtemp()
        for i in range(40):
            im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1, channel_dim=-1)
            n = nib.Nifti1Image(im, np.eye(4))
            nib.save(n, os.path.join(self.data_dir, f"img{i:d}.nii.gz"))
            n = nib.Nifti1Image(seg, np.eye(4))
            nib.save(n, os.path.join(self.data_dir, f"seg{i:d}.nii.gz"))

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0") 

def test_shape(self, img_transform, label_transform, indexes, expected_shape):
        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4))
        tempdir = tempfile.mkdtemp()
        test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
        test_seg1 = os.path.join(tempdir, "test_seg1.nii.gz")
        test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
        test_seg2 = os.path.join(tempdir, "test_seg2.nii.gz")
        nib.save(test_image, test_image1)
        nib.save(test_image, test_seg1)
        nib.save(test_image, test_image2)
        nib.save(test_image, test_seg2)
        test_images = [test_image1, test_image2]
        test_segs = [test_seg1, test_seg2]
        test_labels = [1, 1]
        dataset = ArrayDataset(test_images, img_transform, test_segs, label_transform, test_labels, None)
        self.assertEqual(len(dataset), 2)
        dataset.set_random_state(1234)
        data1 = dataset[0]
        data2 = dataset[1]

        self.assertTupleEqual(data1[indexes[0]].shape, expected_shape)
        self.assertTupleEqual(data1[indexes[1]].shape, expected_shape)
        np.testing.assert_allclose(data1[indexes[0]], data1[indexes[1]])
        self.assertTupleEqual(data2[indexes[0]].shape, expected_shape)
        self.assertTupleEqual(data2[indexes[1]].shape, expected_shape)
        np.testing.assert_allclose(data2[indexes[0]], data2[indexes[0]])

        dataset = ArrayDataset(test_images, img_transform, test_segs, label_transform, test_labels, None)
        dataset.set_random_state(1234)
        _ = dataset[0]
        data2_new = dataset[1]
        np.testing.assert_allclose(data2[indexes[0]], data2_new[indexes[0]], atol=1e-3)
        shutil.rmtree(tempdir) 

def saveScan(img_np, affine, save_path, header = None):
	if header:
		nft_img = nib.Nifti1Image(img_np, affine, header = header)
	else:
		nft_img = nib.Nifti1Image(img_np, affine)
	nib.save(nft_img, save_path)

#get list of validation/train sets 

def register_images(moving: np.ndarray, fixed: np.ndarray, transform_type: str = 'a', n_threads: int = 1) -> np.ndarray:
    """
    Apply RegistrationSynQuick to the input images.

    Parameters
    ----------
    moving: np.ndarray
    fixed: np.ndarray
    transform_type: str, optional
         |  t:  translation
         |  r:  rigid
         |  a:  rigid + affine (default)
         |  s:  rigid + affine + deformable syn
         |  sr: rigid + deformable syn
         |  b:  rigid + affine + deformable b-spline syn
         |  br: rigid + deformable b-spline syn
    n_threads: int, optional
        the number of threads used to apply the registration
    """
    with tempfile.TemporaryDirectory() as tempdir:
        template_path = jp(tempdir, 'template.nii.gz')
        moving_path = jp(tempdir, 'moving.nii.gz')
        nib.save(nib.Nifti1Image(fixed, np.eye(4)), template_path)
        nib.save(nib.Nifti1Image(moving, np.eye(4)), moving_path)

        reg = RegistrationSynQuick()
        reg.inputs.fixed_image = template_path
        reg.inputs.moving_image = moving_path
        reg.inputs.num_threads = n_threads
        reg.inputs.transform_type = transform_type
        reg.inputs.output_prefix = jp(tempdir, 'transform')
        reg.run()

        return nib.load(jp(tempdir, 'transformWarped.nii.gz')).get_data() 

def test_shape(self, expected_shape):
        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
        tempdir = tempfile.mkdtemp()
        nib.save(test_image, os.path.join(tempdir, "test_image1.nii.gz"))
        nib.save(test_image, os.path.join(tempdir, "test_label1.nii.gz"))
        nib.save(test_image, os.path.join(tempdir, "test_extra1.nii.gz"))
        nib.save(test_image, os.path.join(tempdir, "test_image2.nii.gz"))
        nib.save(test_image, os.path.join(tempdir, "test_label2.nii.gz"))
        nib.save(test_image, os.path.join(tempdir, "test_extra2.nii.gz"))
        test_data = [
            {
                "image": os.path.join(tempdir, "test_image1.nii.gz"),
                "label": os.path.join(tempdir, "test_label1.nii.gz"),
                "extra": os.path.join(tempdir, "test_extra1.nii.gz"),
            },
            {
                "image": os.path.join(tempdir, "test_image2.nii.gz"),
                "label": os.path.join(tempdir, "test_label2.nii.gz"),
                "extra": os.path.join(tempdir, "test_extra2.nii.gz"),
            },
        ]
        dataset = CacheDataset(
            data=test_data, transform=Compose([LoadNiftid(keys=["image", "label", "extra"])]), cache_rate=0.5
        )
        data1 = dataset[0]
        data2 = dataset[1]
        shutil.rmtree(tempdir)
        self.assertTupleEqual(data1["image"].shape, expected_shape)
        self.assertTupleEqual(data1["label"].shape, expected_shape)
        self.assertTupleEqual(data1["extra"].shape, expected_shape)
        self.assertTupleEqual(data2["image"].shape, expected_shape)
        self.assertTupleEqual(data2["label"].shape, expected_shape)
        self.assertTupleEqual(data2["extra"].shape, expected_shape) 

def save_landmarks(path_file, landmarks):
    """ save landmarks into a specific file

    both used formats csv/pts is using the same coordinate frame,
    the origin (0, 0) is located in top left corner of the image

    :param str path_file: path to the output file
    :param landmarks: np.array<np_points, dim>
    """
    assert os.path.isdir(os.path.dirname(path_file)), \
        'missing folder "%s"' % os.path.dirname(path_file)
    path_file, _ = os.path.splitext(path_file)
    landmarks = landmarks.values if isinstance(landmarks, pd.DataFrame) else landmarks
    save_landmarks_csv(path_file + '.csv', landmarks)
    save_landmarks_pts(path_file + '.pts', landmarks) 

def save_numpy(value, path: PathLike, *, allow_pickle: bool = True, fix_imports: bool = True):
    """A wrapper around ``np.save`` that matches the interface ``save(what, where)``."""
    np.save(path, value, allow_pickle=allow_pickle, fix_imports=fix_imports) 

def test_shape(self, input_param, expected_shape):
        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
        test_data = dict()
        tempdir = tempfile.mkdtemp()
        for key in KEYS:
            nib.save(test_image, os.path.join(tempdir, key + ".nii.gz"))
            test_data.update({key: os.path.join(tempdir, key + ".nii.gz")})
        result = LoadNiftid(**input_param)(test_data)
        for key in KEYS:
            self.assertTupleEqual(result[key].shape, expected_shape)
        shutil.rmtree(tempdir)