diff --git a/brainiak/fcma/preprocessing.py b/brainiak/fcma/preprocessing.py index e9f1b302d..5a620237d 100644 --- a/brainiak/fcma/preprocessing.py +++ b/brainiak/fcma/preprocessing.py @@ -23,7 +23,7 @@ from scipy.stats.mstats import zscore from mpi4py import MPI -from ..image import multimask_images +from ..image import mask_images, multimask_images logger = logging.getLogger(__name__) @@ -123,12 +123,13 @@ def prepare_fcma_data(images, conditions, mask1, mask2=None, raw_data2 = [] if rank == 0: if mask2 is not None: - activity_data1, activity_data2 = multimask_images(images, - [mask1, mask2], - np.float32) + masks = (mask1, mask2) + activity_data1, activity_data2 = zip(*multimask_images(images, + masks, + np.float32)) raw_data2, _ = _separate_epochs(activity_data2, conditions) else: - (activity_data1,) = multimask_images(images, (mask1,), np.float32) + activity_data1 = list(mask_images(images, mask1, np.float32)) raw_data1, labels = _separate_epochs(activity_data1, conditions) time1 = time.time() raw_data_length = len(raw_data1) @@ -216,7 +217,7 @@ def prepare_mvpa_data(images, conditions, mask): labels: 1D array contains labels of the data """ - (activity_data,) = multimask_images(images, (mask,), np.float32) + activity_data = list(mask_images(images, mask, np.float32)) epoch_info = generate_epochs_info(conditions) num_epochs = len(epoch_info) (d1, _) = activity_data[0].shape diff --git a/brainiak/image.py b/brainiak/image.py index fc0d85e5c..91c6c2d40 100644 --- a/brainiak/image.py +++ b/brainiak/image.py @@ -19,14 +19,63 @@ "mask_image", "multimask_images", ] -from typing import Iterable, Sequence -from typing import List # noqa F401 https://gitlab.com/pycqa/flake8/issues/118 + +import itertools + +from typing import Iterable, Sequence, Type, TypeVar import numpy as np from nibabel.spatialimages import SpatialImage +T = TypeVar("T", bound="MaskedMultiSubjectData") + + +class MaskedMultiSubjectData(np.ndarray): + """Array in shape n_voxels, n_trs, n_subjects.""" + @classmethod + def from_masked_images(cls: Type[T], masked_images: Iterable[np.ndarray], + n_sub: int) -> T: + """Create a new instance from masked images. + + Parameters + ---------- + masked_images + Images to concatenate. + n_sub + Number of subjects. Must match the number of images. + + Returns + ------- + T + A new instance. + + Raises + ------ + ValueError + Images have different shapes. + + The number of images differs from n_sub. + """ + images_iterator = iter(masked_images) + first_image = next(images_iterator) + result = np.empty((first_image.shape[0], first_image.shape[1], n_sub)) + for n_images, image in enumerate(itertools.chain([first_image], + images_iterator)): + if image.shape != first_image.shape: + raise ValueError("Image {} has different shape from first " + "image: {} != {}".format(n_images, + image.shape, + first_image.shape)) + result[:, :, n_images] = image + n_images += 1 + if n_images != n_sub: + raise ValueError("n_sub != number of images: {} != {}" + .format(n_sub, n_images)) + return result.view(cls) + + class ConditionSpec(np.ndarray): """One-hot representation of conditions across epochs and TRs. @@ -85,7 +134,7 @@ def mask_image(image: SpatialImage, mask: np.ndarray, data_type: type = None def multimask_images(images: Iterable[SpatialImage], masks: Sequence[np.ndarray], image_type: type = None - ) -> Sequence[Sequence[np.ndarray]]: + ) -> Iterable[Sequence[np.ndarray]]: """Mask images with multiple masks. Parameters @@ -97,14 +146,32 @@ def multimask_images(images: Iterable[SpatialImage], image_type: Type to cast images to. - Returns - ------- - List[List[np.ndarray]] - For each mask, a list of masked images. + Yields + ------ + Sequence[np.ndarray] + For each mask, a masked image. """ - masked_images = [[] for _ in range(len(masks)) - ] # type: List[List[np.ndarray]] for image in images: - for i, mask in enumerate(masks): - masked_images[i].append(mask_image(image, mask, image_type)) - return masked_images + yield [mask_image(image, mask, image_type) for mask in masks] + + +def mask_images(images: Iterable[SpatialImage], mask: np.ndarray, + image_type: type = None) -> Iterable[np.ndarray]: + """Mask images. + + Parameters + ---------- + images: + Images to mask. + mask: + Mask to apply. + image_type: + Type to cast images to. + + Yields + ------ + np.ndarray + Masked image. + """ + for images in multimask_images(images, (mask,), image_type): + yield images[0] diff --git a/brainiak/io.py b/brainiak/io.py index 00aae5a7d..62937bdae 100644 --- a/brainiak/io.py +++ b/brainiak/io.py @@ -21,7 +21,7 @@ ] from pathlib import Path -from typing import Iterable, List, Union +from typing import Callable, Iterable, List, Union import nibabel as nib import numpy as np @@ -61,13 +61,46 @@ def load_images_from_dir(in_dir: Union[str, Path], suffix: str = "nii.gz", yield nib.load(str(f)) -def load_boolean_mask(path: Union[str, Path]) -> np.ndarray: +def load_images(image_paths: Iterable[Union[str, Path]] + ) -> Iterable[SpatialImage]: + """Load images from paths. + + For efficiency, returns an iterator, not a sequence, so the results cannot + be accessed by indexing. + + For every new iteration through the images, load_images must be called + again. + + Parameters + ---------- + image_paths: + Paths to images. + + Yields + ------ + SpatialImage + Image. + """ + for image_path in image_paths: + if isinstance(image_path, Path): + string_path = str(image_path) + else: + string_path = image_path + yield nib.load(string_path) + + +def load_boolean_mask(path: Union[str, Path], + predicate: Callable[[np.ndarray], np.ndarray] = None + ) -> np.ndarray: """Load boolean nibabel.SpatialImage mask. Parameters ---------- path Mask path. + predicate + Callable used to create boolean values, e.g. a threshold function + ``lambda x: x > 50``. Returns ------- @@ -76,7 +109,12 @@ def load_boolean_mask(path: Union[str, Path]) -> np.ndarray: """ if not isinstance(path, str): path = str(path) - return nib.load(path).get_data().astype(np.bool) + data = nib.load(path).get_data() + if predicate is not None: + mask = predicate(data) + else: + mask = data.astype(np.bool) + return mask def load_labels(path: Union[str, Path]) -> List[SingleConditionSpec]: diff --git a/brainiak/isfc.py b/brainiak/isfc.py index dca2a8670..176042d8b 100644 --- a/brainiak/isfc.py +++ b/brainiak/isfc.py @@ -31,7 +31,6 @@ # Princeton University, 2017 from brainiak.fcma.util import compute_correlation -import nibabel as nib import numpy as np from scipy import stats @@ -115,53 +114,3 @@ def isfc(D, collapse_subj=True): if collapse_subj: ISFC = np.mean(ISFC, axis=2) return ISFC - - -def load_subjects_nii(data_files, mask_file, mask_func=None): - """Loading masked nifti data into matrix - - Given a list of subject data files and a mask, loads voxel timecourses - inside mask into a matrix, which is returned along with the voxel - coordinates of the mask. If mask_func is given, it specifies which mask - values should be included (e.g. for thresholding a continuous-valued mask) - - Parameters - ---------- - data_files : list of filenames of subject nii files - each nii file should be 4D (space+time), with all dimensions identical - for all subjects - - mask_file : filename of mask nii - mask should be 3D (space), with dimensions identical to the first - three dimensions of the data files - - mask_func : Callable[[ndarray], bool] : default x>0 - - Returns - ------- - D : voxel by time by subject ndarray - all data within mask, from all subjects - coords : tuple of 3 ndarrays - x,y,z (as provided by nibabel) coordinates of mask voxel locations - """ - - mask_nii = nib.load(mask_file) - mask = mask_nii.get_data() - if mask_func is not None: - mask = mask_func(mask) - else: - mask = mask > 0 - mask_shape = mask_nii.shape - coords = np.where(mask) - - data_shape = nib.load(data_files[0]).shape - D = np.zeros((np.sum(mask), data_shape[3], len(data_files))) - for s in range(len(data_files)): - nii = nib.load(data_files[s]) - if nii.shape != data_shape: - raise ValueError("Data has different shapes across subjects") - if nii.shape[:3] != mask_shape: - raise ValueError("Data and mask have different shapes") - D[:, :, s] = nii.get_data()[mask] - - return (D, coords) diff --git a/examples/isfc/isfc.py b/examples/isfc/isfc.py index 9d0bfd2f1..373755623 100644 --- a/examples/isfc/isfc.py +++ b/examples/isfc/isfc.py @@ -21,6 +21,7 @@ # Princeton University, 2017 import brainiak.isfc +from brainiak import image, io import nibabel as nib import numpy as np from matplotlib import pyplot as plt @@ -36,8 +37,11 @@ subj in np.arange(1, 5)] print('Loading data from ', len(fnames), ' subjects...') -D, coords = brainiak.isfc.load_subjects_nii(fnames, brain_fname, - lambda x: x > 50) + +brain_mask = io.load_boolean_mask(brain_fname, lambda x: x > 50) +masked_images = image.mask_images(io.load_images(fnames), brain_mask) +coords = np.where(brain_mask) +D = image.MaskedMultiSubjectData.from_masked_images(masked_images, len(fnames)) print('Calculating ISC on ', D.shape[0], ' voxels') ISC = brainiak.isfc.isc(D) diff --git a/tests/image/test_image.py b/tests/image/test_image.py index 79f49565a..7c56f139f 100644 --- a/tests/image/test_image.py +++ b/tests/image/test_image.py @@ -20,8 +20,21 @@ from nibabel.nifti1 import Nifti1Pair from nibabel.spatialimages import SpatialImage -from brainiak.image import (mask_image, multimask_images, - SingleConditionSpec) +from brainiak.image import (mask_image, mask_images, MaskedMultiSubjectData, + multimask_images, SingleConditionSpec) + + +@pytest.fixture +def masked_multi_subject_data(masked_images): + return np.stack(masked_images, axis=-1) + + +class TestMaskedMultiSubjectData: + def test_from_masked_images(self, masked_images, + masked_multi_subject_data): + result = MaskedMultiSubjectData.from_masked_images(masked_images, + len(masked_images)) + assert np.array_equal(result, masked_multi_subject_data) @pytest.fixture @@ -42,22 +55,22 @@ def test_extract_labels(self, condition_spec: SingleConditionSpec @pytest.fixture def spatial_image() -> SpatialImage: - return Nifti1Pair(np.array([[[0, 0, 0, 0], - [0, 0, 0, 0], - [0, 0, 0, 0], - [0, 0, 0, 0]], - [[0, 0, 0, 0], - [0, 1, 0, 0], - [0, 0, 1, 0], - [0, 0, 0, 0]], - [[0, 0, 0, 0], - [0, 0, 1, 0], - [0, 1, 0, 0], - [0, 0, 0, 0]], - [[0, 0, 0, 0], - [0, 0, 0, 0], - [0, 0, 0, 0], - [0, 0, 0, 0]]]), + return Nifti1Pair(np.array([[[[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], + [[0, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 0]], + [[0, 0, 0, 0], + [0, 0, 1, 0], + [0, 1, 0, 0], + [0, 0, 0, 0]], + [[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]]]]).reshape(4, 4, 4, 1), np.eye(4)) @@ -83,14 +96,14 @@ def mask() -> np.ndarray: @pytest.fixture def masked_data() -> np.ndarray: - return np.array([1, 0, 0, 1, 0, 1, 1, 0]) + return np.array([[1, 0, 0, 1, 0, 1, 1, 0]]).reshape(8, 1) @pytest.fixture def images(spatial_image: SpatialImage) -> Iterable[SpatialImage]: images = [spatial_image] image_data = spatial_image.get_data().copy() - image_data[1, 1, 1] = 2 + image_data[1, 1, 1, 0] = 2 images.append(Nifti1Pair(image_data, np.eye(4))) return images @@ -108,11 +121,17 @@ def masks(mask: np.ndarray) -> Sequence[np.ndarray]: @pytest.fixture -def multimasked_data(masked_data) -> Iterable[Iterable[np.ndarray]]: - masked_data_2 = np.hstack((2, masked_data[1:])) - return [[masked_data, masked_data_2], - [np.hstack((0, masked_data)), np.hstack((0, masked_data_2))], - [masked_data[:-1], masked_data_2[:-1]]] +def multimasked_images(masked_data) -> Iterable[Iterable[np.ndarray]]: + masked_data_2 = np.concatenate(([[2]], masked_data[1:, :])) + return [[masked_data, np.concatenate(([[0]], masked_data)), + masked_data[:-1, :]], + [masked_data_2, np.concatenate(([[0]], masked_data_2)), + masked_data_2[:-1, :]]] + + +@pytest.fixture +def masked_images(multimasked_images) -> Iterable[np.ndarray]: + return [multimasked_image[0] for multimasked_image in multimasked_images] def test_mask_image(spatial_image: SpatialImage, mask: np.ndarray, @@ -129,11 +148,24 @@ def test_mask_image_with_type(spatial_image: SpatialImage, mask: np.ndarray, assert np.allclose(result, masked_data) -def test_multimask_images(images: Iterable[SpatialImage], - masks: Sequence[np.ndarray], - multimasked_data: Iterable[Iterable[np.ndarray]] - ) -> None: +def test_multimask_images( + images: Iterable[SpatialImage], + masks: Sequence[np.ndarray], + multimasked_images: Iterable[Iterable[np.ndarray]] + ) -> None: result = multimask_images(images, masks) - for mask_data in zip(result, multimasked_data): - for result_data, precomputed_data in zip(mask_data[0], mask_data[1]): - assert np.array_equal(result_data, precomputed_data) + for result_images, expected_images in zip(result, + multimasked_images): + for result_image, expected_image in zip(result_images, + expected_images): + assert np.array_equal(result_image, expected_image) + + +def test_mask_images( + images: Iterable[SpatialImage], + mask: np.ndarray, + masked_images: Iterable[np.ndarray] + ) -> None: + result = mask_images(images, mask) + for result_image, expected_image in zip(result, masked_images): + assert np.array_equal(result_image, expected_image) diff --git a/tests/io/test_io.py b/tests/io/test_io.py index f00e421e2..9f72550f6 100644 --- a/tests/io/test_io.py +++ b/tests/io/test_io.py @@ -13,7 +13,7 @@ # limitations under the License. from pathlib import Path -from typing import Sequence +from typing import Iterable, Sequence import nibabel as nib import numpy as np @@ -52,16 +52,29 @@ def expected_n_subjects() -> int: return 2 +@pytest.fixture +def image_paths(in_dir: Path) -> Iterable[Path]: + return (in_dir / "subject1_bet.nii.gz", in_dir / "subject2_bet.nii.gz") + + def test_load_images_from_dir_data_shape( in_dir: Path, expected_image_data_shape: Sequence[int], expected_n_subjects: int ) -> None: - i = 0 - for image in io.load_images_from_dir(in_dir, "bet.nii.gz"): + for i, image in enumerate(io.load_images_from_dir(in_dir, "bet.nii.gz")): assert image.get_data().shape == (64, 64, 26, 10) - i += 1 - assert i == expected_n_subjects + assert i + 1 == expected_n_subjects + + +def test_load_images_data_shape( + image_paths: Iterable[Path], + expected_image_data_shape: Sequence[int], + expected_n_subjects: int + ) -> None: + for i, image in enumerate(io.load_images(image_paths)): + assert image.get_data().shape == (64, 64, 26, 10) + assert i + 1 == expected_n_subjects def test_load_boolean_mask(mask_path: Path) -> None: @@ -69,6 +82,12 @@ def test_load_boolean_mask(mask_path: Path) -> None: assert mask.dtype == np.bool +def test_load_boolean_mask_predicate(mask_path: Path) -> None: + mask = io.load_boolean_mask(mask_path, lambda x: np.logical_not(x)) + expected_mask = np.logical_not(io.load_boolean_mask(mask_path)) + assert np.array_equal(mask, expected_mask) + + def test_load_labels(labels_path: Path, expected_condition_spec_shape: Sequence[int], expected_n_subjects: int) -> None: diff --git a/tests/isfc/test_isfc.py b/tests/isfc/test_isfc.py index c16d66b8d..a3bfb53f1 100644 --- a/tests/isfc/test_isfc.py +++ b/tests/isfc/test_isfc.py @@ -1,4 +1,5 @@ import brainiak.isfc +from brainiak import image, io import numpy as np import sys, os @@ -18,14 +19,17 @@ def test_ISC(): assert np.isclose(ISC, [0.9540602, 0.99585304]).all(), \ "Calculated ISC does not match ground truth" -def test_loading_and_ISFC(): +def test_ISFC(): curr_dir = os.path.dirname(__file__) mask_fname = os.path.join(curr_dir,'mask.nii.gz') + mask = io.load_boolean_mask(mask_fname) fnames = [os.path.join(curr_dir,'subj1.nii.gz'), os.path.join(curr_dir,'subj2.nii.gz')] + masked_images = image.mask_images(io.load_images(fnames), mask) - D, coords = brainiak.isfc.load_subjects_nii(fnames, mask_fname) + D = image.MaskedMultiSubjectData.from_masked_images(masked_images, + len(fnames)) assert D.shape == (4,5,2), "Loaded data has incorrect shape"