Repository navigation
Expand file tree
/
Copy pathtest_io.py
More file actions
106 lines (77 loc) · 3.01 KB
/
Copy pathtest_io.py
File metadata and controls
106 lines (77 loc) · 3.01 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
# Copyright 2017 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from pathlib import Path
from typing import Iterable, Sequence
import nibabel as nib
import numpy as np
import pytest
from brainiak import io
@pytest.fixture
def in_dir() -> Path:
return Path(__file__).parent / "data"
@pytest.fixture
def expected_image_data_shape() -> Sequence[int]:
return (64, 64, 26, 10)
@pytest.fixture
def mask_path(in_dir: Path) -> Path:
return in_dir / "mask.nii.gz"
@pytest.fixture
def labels_path(in_dir: Path) -> Path:
return in_dir / "epoch_labels.npy"
@pytest.fixture
def expected_condition_spec_shape() -> Sequence[int]:
return (2, 2, 10)
@pytest.fixture
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:
for i, image in enumerate(io.load_images_from_dir(in_dir, "bet.nii.gz")):
assert image.get_data().shape == (64, 64, 26, 10)
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:
mask = io.load_boolean_mask(mask_path)
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:
condition_specs = io.load_labels(labels_path)
i = 0
for condition_spec in condition_specs:
assert condition_spec.shape == expected_condition_spec_shape
i += 1
assert i == expected_n_subjects
def test_save_as_nifti_file(tmpdir) -> None:
out_file = str(tmpdir / "nifti.nii")
shape = (4, 4, 4)
io.save_as_nifti_file(np.ones(shape), np.eye(4), out_file)
assert nib.load(out_file).get_data().shape == shape