Python nibabel.Nifti1Image() 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 betaseries_file(tmpdir_factory,
                    deriv_betaseries_fname=deriv_betaseries_fname):
    bfile = tmpdir_factory.mktemp("beta").ensure(deriv_betaseries_fname)
    np.random.seed(3)
    num_trials = 40
    tgt_corr = 0.1
    bs1 = np.random.rand(num_trials)
    # create another betaseries with a target correlation
    bs2 = minimize(lambda x: abs(tgt_corr - pearsonr(bs1, x)[0]),
                   np.random.rand(num_trials)).x

    # two identical beta series
    bs_data = np.array([[[bs1, bs2]]])

    # the nifti image
    bs_img = nib.Nifti1Image(bs_data, np.eye(4))
    bs_img.to_filename(str(bfile))

    return bfile 

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 _dti(self, row):
        dti_params_file = self._get_fname(row, '_model-DTI_diffmodel.nii.gz')
        if self.force_recompute or not op.exists(dti_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()
            dtf = dti_fit(gtab, data, mask=mask)
            self.log_and_save_nii(nib.Nifti1Image(dtf.model_params,
                                                  row['dwi_affine']),
                                  dti_params_file)
            meta_fname = self._get_fname(row, '_model-DTI_diffmodel.json')
            meta = dict(
                Parameters=dict(
                    FitMethod="WLS"),
                OutlierRejection=False,
                ModelURL=f"{DIPY_GH}reconst/dti.py")
            afd.write_json(meta_fname, meta)
        return dti_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 test_dataloading(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)
        loader = DataLoader(dataset, batch_size=10, num_workers=1)
        imgs = next(iter(loader))  # test batching
        np.testing.assert_allclose(imgs.shape, [2] + list(expected_shape))

        dataset.set_random_state(1234)
        new_imgs = next(iter(loader))  # test batching
        np.testing.assert_allclose(imgs, new_imgs, atol=1e-3) 

def read_callosum_templates(resample_to=False):
    """Load AFQ callosum templates from file

    Returns
    -------
    dict with: keys: names of template ROIs and values: nibabel Nifti1Image
    objects from each of the ROI nifti files.
    """
    files, folder = fetch_callosum_templates()
    template_dict = {}
    for f in files:
        img = nib.load(op.join(folder, f))
        if resample_to:
            if isinstance(resample_to, str):
                resample_to = nib.load(resample_to)
            img = nib.Nifti1Image(reg.resample(img.get_fdata(),
                                               resample_to,
                                               img.affine,
                                               resample_to.affine),
                                  resample_to.affine)
        template_dict[f.split('.')[0]] = img
    return template_dict 

def test_calc_dAdt_nifti(self, sphere3_msh):
        affine = np.array([[5., 0., 0., -300],
                           [0., 5., 0., -200],
                           [0., 0., 5., 0.],
                           [0., 0., 0., 1]])
        field = np.ones((121, 81, 41, 3))
        field[..., 1] = 2
        field[..., 2] = 3
        img = nib.Nifti1Image(field, affine)
        coil_matrix = np.array([[0., 1., 0., 0],
                                [1., 0., 0., 0],
                                [0., 0., 1., -100.],
                                [0., 0., 0., 1]])
        dadt = coil._calculate_dadt_nifti(sphere3_msh, img, coil_matrix, 1e6, None)
        sphere3_msh.nodedata.append(dadt)
        assert np.allclose(dadt.value[:, 0], 2e6, atol=1e-6)
        assert np.allclose(dadt.value[:, 1], 1e6, atol=1e-6)
        assert np.allclose(dadt.value[:, 2], 3e6, atol=1e-6) 

def read_stanford_hardi_tractography():
    """
    Reads a minimal tractography from the Stanford dataset.
    """
    files, folder = fetch_stanford_hardi_tractography()
    files_dict = {}
    files_dict['mapping.nii.gz'] = nib.load(
        op.join(afq_home,
                'stanford_hardi_tractography',
                'mapping.nii.gz'))

    files_dict['tractography_subsampled.trk'] = load_trk(
        op.join(afq_home,
                'stanford_hardi_tractography',
                'tractography_subsampled.trk'),
        nib.Nifti1Image(np.zeros((10, 10, 10)), np.eye(4)),
        bbox_valid_check=False,
        trk_header_check=False).streamlines

    return files_dict 

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 s3fs_nifti_write(img, fname, fs=None):
    """
    Write a nifti file straight to S3

    Paramters
    ---------
    img : nib.Nifti1Image class instance
        The image containing data to be written into S3
    fname : string
        Full path (including bucket name and extension) to the S3 location
        where the file is to be saved.
    fs : an s3fs.S3FileSystem class instance, optional
        A file-system to refer to. Default to create a new file-system
    """
    if fs is None:
        fs = s3fs.S3FileSystem()

    bio = BytesIO()
    file_map = img.make_file_map({'image': bio, 'header': bio})
    img.to_file_map(file_map)
    data = gzip.compress(bio.getvalue())
    with fs.open(fname, 'wb') as ff:
        ff.write(data) 

def niwrite(data,affine, name , header=None):
	data[np.isnan(data)]=0
	stdout.write(" + Writing file: %s ...." % name) 
	
	thishead = header
	if thishead == None:
		thishead = head.copy()
		thishead.set_data_shape(list(data.shape))

	outni = nib.Nifti1Image(data,affine,header=thishead)
	outni.set_data_dtype('float64')
	outni.to_filename(name)
	

	print 'done.'

	return outni 

def _brain_mask(self, row, median_radius=4, numpass=1, autocrop=False,
                    vol_idx=None, dilate=10):
        brain_mask_file = self._get_fname(row, '_brain_mask.nii.gz')
        if self.force_recompute or not op.exists(brain_mask_file):
            b0_file = self._b0(row)
            mean_b0_img = nib.load(b0_file)
            mean_b0 = mean_b0_img.get_fdata()
            _, brain_mask = median_otsu(mean_b0, median_radius, numpass,
                                        autocrop, dilate=dilate)
            be_img = nib.Nifti1Image(brain_mask.astype(int),
                                     mean_b0_img.affine)
            self.log_and_save_nii(be_img, brain_mask_file)
            meta = dict(source=b0_file,
                        median_radius=median_radius,
                        numpass=numpass,
                        autocrop=autocrop,
                        vol_idx=vol_idx)
            meta_fname = self._get_fname(row, '_brain_mask.json')
            afd.write_json(meta_fname, meta)
        return brain_mask_file 

def mosaic_to_nii(dcm_data):
    ''' Get Nifti file from Siemens

    Parameters
    ----------
    dcm_data : ``dicom.DataSet``
       DICOM header / image as read by ``dicom`` package

    Returns
    -------
    img : ``Nifti1Image``
       Nifti image object
    '''
    import nibabel as nib
    dcm_w = wrapper_from_data(dcm_data)
    if not dcm_w.is_mosaic:
        raise DicomReadError('data does not appear to be in mosaic format')
    data = dcm_w.get_data()
    aff = np.dot(DPCS_TO_TAL, dcm_w.get_affine())
    return nib.Nifti1Image(data, aff) 

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 _csd(self, row, response=None, sh_order=None, lambda_=1, tau=0.1,):
        csd_params_file = self._get_fname(row, '_model-CSD_diffmodel.nii.gz')
        if self.force_recompute or not op.exists(csd_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()
            csdf = csd_fit(gtab, data, mask=mask,
                           response=response, sh_order=sh_order,
                           lambda_=lambda_, tau=tau)
            self.log_and_save_nii(nib.Nifti1Image(csdf.shm_coeff,
                                                  row['dwi_affine']),
                                  csd_params_file)
            meta_fname = self._get_fname(row, '_model-CSD_diffmodel.json')
            meta = dict(SphericalHarmonicDegree=sh_order,
                        ResponseFunctionTensor=response,
                        SphericalHarmonicBasis="DESCOTEAUX",
                        ModelURL=f"{DIPY_GH}reconst/csdeconv.py",
                        lambda_=lambda_,
                        tau=tau)
            afd.write_json(meta_fname, meta)
        return csd_params_file 

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_load_spacingd_rotate(self, filename):
        data = {"image": filename}
        data_dict = LoadNiftid(keys="image")(data)
        data_dict = AddChanneld(keys="image")(data_dict)
        affine = data_dict["image_meta_dict"]["affine"]
        data_dict["image_meta_dict"]["original_affine"] = data_dict["image_meta_dict"]["affine"] = (
            np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]]) @ affine
        )
        t = time.time()
        res_dict = Spacingd(keys="image", pixdim=(1, 2, 3), diagonal=True, padding_mode="zeros")(data_dict)
        t1 = time.time()
        print(f"time monai: {t1 - t}")
        anat = nibabel.Nifti1Image(data_dict["image"][0], data_dict["image_meta_dict"]["original_affine"])
        ref = resample_to_output(anat, (1, 2, 3), order=1)
        t2 = time.time()
        print(f"time scipy: {t2 - t1}")
        self.assertTrue(t2 >= t1)
        np.testing.assert_allclose(res_dict["image_meta_dict"]["affine"], ref.affine)
        if "anatomical" not in filename:
            np.testing.assert_allclose(res_dict["image"].shape[1:], ref.shape)
            np.testing.assert_allclose(ref.get_fdata(), res_dict["image"][0], atol=0.05)
        else:
            # different from the ref implementation (shape computed by round
            # instead of ceil)
            np.testing.assert_allclose(ref.get_fdata()[..., :-1], res_dict["image"][0], atol=0.05) 

def test_get_clusters_table():
    shape = (9, 10, 11)
    data = np.zeros(shape)
    data[2:4, 5:7, 6:8] = 5.
    stat_img = nib.Nifti1Image(data, np.eye(4))

    # test one cluster extracted
    cluster_table = get_clusters_table(stat_img, 4, 0)
    assert len(cluster_table) == 1

    # test empty table on high stat threshold
    cluster_table = get_clusters_table(stat_img, 6, 0)
    assert len(cluster_table) == 0

    # test empty table on high cluster threshold
    cluster_table = get_clusters_table(stat_img, 4, 9)
    assert len(cluster_table) == 0 

def test_load_spacingd(self, filename):
        data = {"image": filename}
        data_dict = LoadNiftid(keys="image")(data)
        data_dict = AddChanneld(keys="image")(data_dict)
        t = time.time()
        res_dict = Spacingd(keys="image", pixdim=(1, 0.2, 1), diagonal=True, padding_mode="zeros")(data_dict)
        t1 = time.time()
        print(f"time monai: {t1 - t}")
        anat = nibabel.Nifti1Image(data_dict["image"][0], data_dict["image_meta_dict"]["original_affine"])
        ref = resample_to_output(anat, (1, 0.2, 1), order=1)
        t2 = time.time()
        print(f"time scipy: {t2 - t1}")
        self.assertTrue(t2 >= t1)
        np.testing.assert_allclose(res_dict["image_meta_dict"]["affine"], ref.affine)
        np.testing.assert_allclose(res_dict["image"].shape[1:], ref.shape)
        np.testing.assert_allclose(ref.get_fdata(), res_dict["image"][0], atol=0.05) 

def _dki_fa(self, row):
        dki_fa_file = self._get_fname(row, '_model-DKI_FA.nii.gz')
        if self.force_recompute or not op.exists(dki_fa_file):
            tf = self._dki_fit(row)
            fa = tf.fa
            nib.save(nib.Nifti1Image(fa, row['dwi_affine']),
                     dki_fa_file)
            meta_fname = self._get_fname(row, '_model-DKI_FA.json')
            meta = dict()
            afd.write_json(meta_fname, meta)
        return dki_fa_file 

def _dti_pdd(self, row):
        dti_pdd_file = self._get_fname(row, '_model-DTI_PDD.nii.gz')
        if self.force_recompute or not op.exists(dti_pdd_file):
            tf = self._dti_fit(row)
            pdd = tf.directions.squeeze()
            # Invert the x coordinates:
            pdd[..., 0] = pdd[..., 0] * -1

            self.log_and_save_nii(nib.Nifti1Image(pdd, row['dwi_affine']),
                                  dti_pdd_file)
            meta_fname = self._get_fname(row, '_model-DTI_PDD.json')
            meta = dict()
            afd.write_json(meta_fname, meta)
        return dti_pdd_file 

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 _dti_md(self, row):
        dti_md_file = self._get_fname(row, '_model-DTI_MD.nii.gz')
        if self.force_recompute or not op.exists(dti_md_file):
            tf = self._dti_fit(row)
            md = tf.md
            self.log_and_save_nii(nib.Nifti1Image(md, row['dwi_affine']),
                                  dti_md_file)
            meta_fname = self._get_fname(row, '_model-DTI_MD.json')
            meta = dict()
            afd.write_json(meta_fname, meta)
        return dti_md_file 

def _write_fake_bold_img(file_path, shape, rk=3, affine=np.eye(4)):
    data = np.random.randn(*shape)
    data[1:-1, 1:-1, 1:-1] += 100
    Nifti1Image(data, affine).to_filename(file_path)
    return file_path 

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) 

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 test_affine(self):
        test_img = nib.Nifti1Image(np.zeros((1, 2, 3)), np.eye(4) * 20.0)
        test_img = correct_nifti_header_if_necessary(test_img)
        np.testing.assert_allclose(
            test_img.affine,
            np.array([[20.0, 0.0, 0.0, 0.0], [0.0, 20.0, 0.0, 0.0], [0.0, 0.0, 20.0, 0.0], [0.0, 0.0, 0.0, 20.0]]),
        ) 

def save_segmentation(seg, cid, output_path):
    # Resolve location where data should be written
    if not os.path.exists(output_path):
        raise IOError(
            "Data path, {}, could not be resolved".format(str(output_path))
        )
    # Convert numpy array to NIFTI
    nifti = nib.Nifti1Image(seg, None)
    # nifti.get_data_dtype() = seg.dtype
    # Save segmentation to disk
    nib.save(nifti, os.path.join(output_path,
                                 "prediction_" + str(cid).zfill(5) + ".nii.gz")) 

def _write_fake_fmri_data(shapes, rk=3, affine=np.eye(4)):
    mask_file, fmri_files, design_files = 'mask.nii', [], []
    for i, shape in enumerate(shapes):
        fmri_files.append('fmri_run%d.nii' % i)
        data = np.random.randn(*shape)
        data[1:-1, 1:-1, 1:-1] += 100
        Nifti1Image(data, affine).to_filename(fmri_files[-1])
        design_files.append('dmtx_%d.csv' % i)
        pd.DataFrame(np.random.randn(shape[3], rk),
                     columns=['', '', '']).to_csv(design_files[-1])
    Nifti1Image((np.random.rand(*shape[:3]) > .5).astype(np.int8),
                affine).to_filename(mask_file)
    return mask_file, fmri_files, design_files