def test_sample_image_label_one_label(self, shape: Tuple):
"""
Test sample_image_label in labeled case with one label.
:param shape: shape of the label.
"""
loader = GeneratorDataLoader(labeled=True,
num_indices=1,
sample_label="all")
got = next(
loader.sample_image_label(
moving_image=get_arr(shape=shape[:3], seed=0),
fixed_image=get_arr(shape=shape[:3], seed=1),
moving_label=get_arr(shape=shape, seed=2),
fixed_label=get_arr(shape=shape, seed=3),
image_indices=[1],
))
expected = dict(
moving_image=get_arr(shape=shape[:3], seed=0),
fixed_image=get_arr(shape=shape[:3], seed=1),
moving_label=get_arr(shape=shape[:3], seed=2),
fixed_label=get_arr(shape=shape[:3], seed=3),
indices=np.asarray([1, 0], dtype=np.float32),
)
assert all(
is_equal_np(got[key], expected[key]) for key in expected.keys())
def test_sample_image_label_multiple_labels(self):
"""Test sample_image_label in labeled case with multiple labels."""
loader = GeneratorDataLoader(labeled=True,
num_indices=1,
sample_label="all")
shape = (2, 3, 4, 5)
got_iter = loader.sample_image_label(
moving_image=get_arr(shape=shape[:3], seed=0),
fixed_image=get_arr(shape=shape[:3], seed=1),
moving_label=get_arr(shape=shape, seed=2),
fixed_label=get_arr(shape=shape, seed=3),
image_indices=[1],
)
moving_label = get_arr(shape=shape, seed=2)
fixed_label = get_arr(shape=shape, seed=3)
for i in range(shape[-1]):
got = next(got_iter)
expected = dict(
moving_image=get_arr(shape=shape[:3], seed=0),
fixed_image=get_arr(shape=shape[:3], seed=1),
moving_label=moving_label[:, :, :, i],
fixed_label=fixed_label[:, :, :, i],
indices=np.asarray([1, i], dtype=np.float32),
)
assert all(
is_equal_np(got[key], expected[key])
for key in expected.keys())
def test_data_generator(self, labeled: bool):
"""
Test data_generator()
:param labeled: labeled data or not.
"""
class MockDataLoader:
"""Toy data loader."""
def __init__(self, seed: int):
"""
Init.
:param seed: random seed for numpy.
:param kwargs: additional arguments.
"""
self.seed = seed
def get_data(self, index: int) -> np.ndarray:
"""
Return the dummy array despite of the index.
:param index: not used
:return: dummy array.
"""
assert isinstance(index, int)
return get_arr(seed=self.seed)
def mock_sample_index_generator():
"""Toy sample index generator."""
return [[1, 1, [1]]]
loader = GeneratorDataLoader(labeled=labeled,
num_indices=1,
sample_label="all")
loader.__setattr__("sample_index_generator",
mock_sample_index_generator)
loader.loader_moving_image = MockDataLoader(seed=0)
loader.loader_fixed_image = MockDataLoader(seed=1)
if labeled:
loader.loader_moving_label = MockDataLoader(seed=2)
loader.loader_fixed_label = MockDataLoader(seed=3)
# check data loader output
got = next(loader.data_generator())
expected = {
"moving_image": normalize_array(get_arr(seed=0)),
"fixed_image": normalize_array(get_arr(seed=1)),
# 0 or -1 is the label index
"indices": np.array([1, 0] if labeled else [1, -1],
dtype=np.float32),
}
if labeled:
expected = {
"moving_label": get_arr(seed=2),
"fixed_label": get_arr(seed=3),
**expected,
}
assert all(
is_equal_np(got[key], expected[key]) for key in expected.keys())
def test_get_labeled_dataset(self, labeled: bool):
"""
Test get_dataset with data loader.
:param labeled: labeled data or not.
"""
sample = {
"moving_image": get_arr(),
"fixed_image": get_arr(),
"indices": [1],
}
if labeled:
sample = {
"moving_label": get_arr(),
"fixed_label": get_arr(),
**sample,
}
def mock_gen():
"""Toy data generator."""
for _ in range(3):
yield sample
loader = GeneratorDataLoader(labeled=labeled,
num_indices=1,
sample_label="all")
loader.__setattr__("data_generator", mock_gen)
dataset = loader.get_dataset()
for got in dataset.as_numpy_iterator():
assert all(
is_equal_np(got[key], sample[key]) for key in sample.keys())
def test_get_data(self, name, index, expected):
loader = get_loader(name)
array = loader.get_data(index)
got = [
np.shape(array),
[np.amax(array), np.amin(array), np.mean(array), np.std(array)],
]
assert got[0] == expected[0]
assert is_equal_np(got[1], expected[1])
loader.close()
def test_gen_transform_params(self, name: str, moving_param_shape: tuple,
fixed_param_shape: tuple):
"""
Check return shapes and moving/fixed params should be different.
:param name: name of the layer
:param moving_param_shape: params shape for moving image/label
:param fixed_param_shape: params shape for fixed image/label
"""
layer = self.build_layer(name)
moving, fixed = layer.gen_transform_params()
assert moving.shape == (self.batch_size, *moving_param_shape)
assert fixed.shape == (self.batch_size, *fixed_param_shape)
assert not is_equal_np(moving, fixed)
def test_overwrite(self, overwrite: bool):
arr1 = np.random.rand(2, 3, 4, 3)
arr2 = arr1 + 1
nifti_file_path = os.path.join(self.save_dir,
self.arr_name + ".nii.gz")
# save arr1
os.makedirs(self.save_dir, exist_ok=True)
nib.save(img=nib.Nifti1Image(arr1, affine=np.eye(4)),
filename=nifti_file_path)
# save arr2 w/o overwrite
save_array(
save_dir=self.save_dir,
arr=arr2,
name=self.arr_name,
normalize=True,
overwrite=overwrite,
)
arr_read = load_nifti_file(file_path=nifti_file_path)
assert is_equal_np(arr2 if overwrite else arr1, arr_read)
def test_sample_image_label_unlabeled(self):
"""Test sample_image_label in unlabeled case."""
loader = GeneratorDataLoader(labeled=False,
num_indices=1,
sample_label="all")
got = next(
loader.sample_image_label(
moving_image=get_arr(seed=0),
fixed_image=get_arr(seed=1),
moving_label=None,
fixed_label=None,
image_indices=[1],
))
expected = dict(
moving_image=get_arr(seed=0),
fixed_image=get_arr(seed=1),
indices=np.asarray([1, -1], dtype=np.float32),
)
assert all(
is_equal_np(got[key], expected[key]) for key in expected.keys())
def test_normalize_array():
"""
Test normalize array by checking the output values
"""
# no v_min, v_max
arr = np.array([0, 1, 2])
got = util.normalize_array(arr=arr)
expected = np.array([0, 0.5, 1])
assert is_equal_np(got, expected)
arr = np.array([-2, 0, 1, 2])
got = util.normalize_array(arr=arr)
expected = np.array([0, 0.5, 0.75, 1])
assert is_equal_np(got, expected)
# no v_min
arr = np.array([0, 1, 2])
got = util.normalize_array(arr=arr, v_max=1)
expected = np.array([0, 1, 1])
assert is_equal_np(got, expected)
arr = np.array([-2, 0, 1, 2])
got = util.normalize_array(arr=arr, v_max=3)
expected = np.array([0, 0.4, 0.6, 0.8])
assert is_equal_np(got, expected)
# no v_max
arr = np.array([0, 1, 2])
got = util.normalize_array(arr=arr, v_min=1)
expected = np.array([0, 0, 1])
assert is_equal_np(got, expected)
arr = np.array([-2, 0, 1, 2])
got = util.normalize_array(arr=arr, v_min=-3)
expected = np.array([0.2, 0.6, 0.8, 1])
assert is_equal_np(got, expected)
def test_generator_data_loader(caplog):
"""
Test the functions in GeneratorDataLoader
:param caplog: used to check warning message.
"""
generator = GeneratorDataLoader(labeled=True,
num_indices=1,
sample_label="all")
# test properties
assert generator.loader_moving_image is None
assert generator.loader_moving_image is None
assert generator.loader_moving_image is None
assert generator.loader_moving_image is None
# not implemented properties / functions
with pytest.raises(NotImplementedError):
generator.sample_index_generator()
# implemented functions
# test get_Dataset
dummy_array = np.random.random(size=(100, 100, 100)).astype(np.float32)
# for unlabeled data
# mock generator
sequence = [
dict(
moving_image=dummy_array,
fixed_image=dummy_array,
moving_label=dummy_array,
fixed_label=dummy_array,
indices=[1],
) for i in range(3)
]
def mock_generator():
for el in sequence:
yield el
# inputs, no error means passed
generator.data_generator = mock_generator
dataset = generator.get_dataset()
# check dataset output
expected = dict(
moving_image=dummy_array,
fixed_image=dummy_array,
moving_label=dummy_array,
fixed_label=dummy_array,
indices=[1],
)
for got in list(dataset.as_numpy_iterator()):
assert all(
is_equal_np(got[key], expected[key]) for key in expected.keys())
# for unlabeled data
generator_unlabeled = GeneratorDataLoader(labeled=False,
num_indices=1,
sample_label="all")
sequence = [
dict(moving_image=dummy_array, fixed_image=dummy_array, indices=[1])
for i in range(3)
]
# inputs, no error means passed
generator_unlabeled.data_generator = mock_generator
dataset = generator_unlabeled.get_dataset()
# check dataset output
expected = dict(moving_image=dummy_array,
fixed_image=dummy_array,
indices=[1])
for got in list(dataset.as_numpy_iterator()):
assert all(
is_equal_np(got[key], expected[key]) for key in expected.keys())
# test data_generator
# create mock data loader and sample index generator
class MockDataLoader:
def __init__(self, **kwargs):
super().__init__(**kwargs)
def get_data(index):
return dummy_array
def mock_sample_index_generator():
return [[[1], [1], [1]]]
generator = GeneratorDataLoader(labeled=True,
num_indices=1,
sample_label="all")
generator.sample_index_generator = mock_sample_index_generator
generator.loader_moving_image = MockDataLoader
generator.loader_fixed_image = MockDataLoader
generator.loader_moving_label = MockDataLoader
generator.loader_fixed_label = MockDataLoader
# check data generator output
got = next(generator.data_generator())
expected = dict(
moving_image=normalize_array(dummy_array),
fixed_image=normalize_array(dummy_array),
moving_label=dummy_array,
fixed_label=dummy_array,
indices=np.asarray([1] + [0], dtype=np.float32),
)
assert all(is_equal_np(got[key], expected[key]) for key in expected.keys())
# test validate_images_and_labels
with pytest.raises(ValueError) as err_info:
generator.validate_images_and_labels(
fixed_image=None,
moving_image=dummy_array,
moving_label=None,
fixed_label=None,
image_indices=[1],
)
assert "moving image and fixed image must not be None" in str(
err_info.value)
with pytest.raises(ValueError) as err_info:
generator.validate_images_and_labels(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=dummy_array,
fixed_label=None,
image_indices=[1],
)
assert "moving label and fixed label must be both None or non-None" in str(
err_info.value)
with pytest.raises(ValueError) as err_info:
generator.validate_images_and_labels(
fixed_image=dummy_array,
moving_image=dummy_array + 1.0,
moving_label=None,
fixed_label=None,
image_indices=[1],
)
assert "Sample [1]'s moving_image's values are not between [0, 1]" in str(
err_info.value)
with pytest.raises(ValueError) as err_info:
generator.validate_images_and_labels(
fixed_image=dummy_array,
moving_image=np.random.random(size=(100, 100)),
moving_label=None,
fixed_label=None,
image_indices=[1],
)
assert "Sample [1]'s moving_image' shape should be 3D. " in str(
err_info.value)
with pytest.raises(ValueError) as err_info:
generator.validate_images_and_labels(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=np.random.random(size=(100, 100)),
fixed_label=dummy_array,
image_indices=[1],
)
assert "Sample [1]'s moving_label' shape should be 3D or 4D. " in str(
err_info.value)
with pytest.raises(ValueError) as err_info:
generator.validate_images_and_labels(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=np.random.random(size=(100, 100, 100, 3)),
fixed_label=np.random.random(size=(100, 100, 100, 4)),
image_indices=[1],
)
assert (
"Sample [1]'s moving image and fixed image have different numbers of labels."
in str(err_info.value))
# warning
caplog.clear() # clear previous log
generator.validate_images_and_labels(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=np.random.random(size=(100, 100, 90)),
fixed_label=dummy_array,
image_indices=[1],
)
assert "Sample [1]'s moving image and label have different shapes. " in caplog.text
caplog.clear() # clear previous log
generator.validate_images_and_labels(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=dummy_array,
fixed_label=np.random.random(size=(100, 100, 90)),
image_indices=[1],
)
assert "Sample [1]'s fixed image and label have different shapes. " in caplog.text
# test sample_image_label method
# for unlabeled input data
got = next(
generator.sample_image_label(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=None,
fixed_label=None,
image_indices=[1],
))
expected = dict(
moving_image=dummy_array,
fixed_image=dummy_array,
indices=np.asarray([1] + [-1], dtype=np.float32),
)
assert all(is_equal_np(got[key], expected[key]) for key in expected.keys())
# for data with one label
got = next(
generator.sample_image_label(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=dummy_array,
fixed_label=dummy_array,
image_indices=[1],
))
expected = dict(
moving_image=dummy_array,
fixed_image=dummy_array,
moving_label=dummy_array,
fixed_label=dummy_array,
indices=np.asarray([1] + [0], dtype=np.float32),
)
assert all(is_equal_np(got[key], expected[key]) for key in expected.keys())
# for data with multiple labels
dummy_labels = np.random.random(size=(100, 100, 100, 3))
got = generator.sample_image_label(
fixed_image=dummy_array,
moving_image=dummy_array,
moving_label=dummy_labels,
fixed_label=dummy_labels,
image_indices=[1],
)
for label_index in range(dummy_labels.shape[3]):
got_iter = next(got)
expected = dict(
moving_image=dummy_array,
fixed_image=dummy_array,
moving_label=dummy_labels[..., label_index],
fixed_label=dummy_labels[..., label_index],
indices=np.asarray([1] + [label_index], dtype=np.float32),
)
assert all(
is_equal_np(got_iter[key], expected[key])
for key in expected.keys())
def test_save_array():
"""
Test save_array by testing different shapes and count output files
"""
def get_num_pngs_in_dir(dir_paths):
return len([x for x in os.listdir(dir_paths) if x.endswith(".png")])
def get_num_niftis_in_dir(dir_paths):
return len([x for x in os.listdir(dir_paths) if x.endswith(".nii.gz")])
save_dir = "logs/test_util_save_array"
if os.path.exists(save_dir):
shutil.rmtree(save_dir)
# test 3D tf tensor
name = "3d_tf"
out_dir = os.path.join(save_dir, name)
arr = tf.random.uniform(shape=(2, 3, 4))
save_array(save_dir=save_dir, arr=arr, name=name, gray=True)
assert get_num_pngs_in_dir(out_dir) == 4
assert get_num_niftis_in_dir(save_dir) == 1
shutil.rmtree(out_dir)
os.remove(os.path.join(save_dir, name + ".nii.gz"))
# test 4D tf tensor
name = "4d_tf"
out_dir = os.path.join(save_dir, name)
arr = tf.random.uniform(shape=(2, 3, 4, 3))
save_array(save_dir=save_dir, arr=arr, name=name, gray=True)
assert get_num_pngs_in_dir(out_dir) == 4
assert get_num_niftis_in_dir(save_dir) == 1
shutil.rmtree(out_dir)
os.remove(os.path.join(save_dir, name + ".nii.gz"))
# test 3D np tensor
name = "3d_np"
out_dir = os.path.join(save_dir, name)
arr = np.random.rand(2, 3, 4)
save_array(save_dir=save_dir, arr=arr, name=name, gray=True)
assert get_num_pngs_in_dir(out_dir) == 4
assert get_num_niftis_in_dir(save_dir) == 1
shutil.rmtree(out_dir)
os.remove(os.path.join(save_dir, name + ".nii.gz"))
# test 4D np tensor
name = "4d_np"
out_dir = os.path.join(save_dir, name)
arr = np.random.rand(2, 3, 4, 3)
save_array(save_dir=save_dir, arr=arr, name=name, gray=True)
assert get_num_pngs_in_dir(out_dir) == 4
assert get_num_niftis_in_dir(save_dir) == 1
shutil.rmtree(out_dir)
os.remove(os.path.join(save_dir, name + ".nii.gz"))
# test 4D np tensor without nifti
name = "4d_np"
out_dir = os.path.join(save_dir, name)
arr = np.random.rand(2, 3, 4, 3)
save_array(save_dir=save_dir,
arr=arr,
name=name,
gray=True,
save_nifti=False)
assert get_num_pngs_in_dir(out_dir) == 4
assert get_num_niftis_in_dir(save_dir) == 0
shutil.rmtree(out_dir)
# test 4D np tensor without png
name = "4d_np"
out_dir = os.path.join(save_dir, name)
arr = np.random.rand(2, 3, 4, 3)
assert not os.path.exists(out_dir)
save_array(save_dir=save_dir,
arr=arr,
name=name,
gray=True,
save_png=False)
assert not os.path.exists(out_dir)
assert get_num_niftis_in_dir(save_dir) == 1
os.remove(os.path.join(save_dir, name + ".nii.gz"))
# test 4D np tensor with overwrite
name = "4d_np"
out_dir = os.path.join(save_dir, name)
arr1 = np.random.rand(2, 3, 4, 3)
arr2 = np.random.rand(2, 3, 4, 3)
assert not is_equal_np(arr1, arr2)
nifti_file_path = os.path.join(save_dir, name + ".nii.gz")
# save arr1
os.makedirs(save_dir, exist_ok=True)
nib.save(img=nib.Nifti2Image(arr1, affine=np.eye(4)),
filename=nifti_file_path)
# save arr2 without overwrite
save_array(save_dir=save_dir,
arr=arr1,
name=name,
gray=True,
overwrite=False)
arr_read = load_nifti_file(file_path=nifti_file_path)
assert is_equal_np(arr1, arr_read)
# save arr2 with overwrite
save_array(save_dir=save_dir,
arr=arr2,
name=name,
gray=True,
overwrite=True)
arr_read = load_nifti_file(file_path=nifti_file_path)
assert is_equal_np(arr2, arr_read)
shutil.rmtree(out_dir)
os.remove(os.path.join(save_dir, name + ".nii.gz"))
# test 5D np tensor
name = "5d_np"
arr = np.random.rand(2, 3, 4, 1, 3)
with pytest.raises(ValueError) as err_info:
save_array(save_dir=save_dir, arr=arr, name=name, gray=True)
assert "arr must be 3d or 4d numpy array or tf tensor" in str(
err_info.value)
# test 4D np tensor with wrong shape
name = "5d_np"
arr = np.random.rand(2, 3, 4, 1)
with pytest.raises(ValueError) as err_info:
save_array(save_dir=save_dir, arr=arr, name=name, gray=True)
assert "4d arr must have 3 channels as last dimension" in str(
err_info.value)
def test_no_max(self, arr, v_min, expected):
got = util.normalize_array(arr=arr, v_min=v_min)
assert is_equal_np(got, expected)
def test_get_data():
"""
test get_data method by verifying returns' shape and value stats
"""
# paired
dir_paths = ["./data/test/nifti/paired/test"]
name = "fixed_images"
loader = NiftiFileLoader(dir_paths=dir_paths, name=name, grouped=False)
index = 0
array = loader.get_data(index)
got = [
np.shape(array),
[np.amax(array),
np.amin(array),
np.mean(array),
np.std(array)],
]
expected = [(44, 59, 41), [255.0, 0.0, 68.359276, 65.84009]]
loader.close()
assert got[0] == expected[0]
assert is_equal_np(got[1], expected[1])
# unpaired
dir_paths = ["./data/test/nifti/unpaired/test"]
name = "images"
loader = NiftiFileLoader(dir_paths=dir_paths, name=name, grouped=False)
index = 0
array = loader.get_data(index)
got = [
np.shape(array),
[np.amax(array),
np.amin(array),
np.mean(array),
np.std(array)],
]
expected = [(64, 64, 60), [255.0, 0.0, 60.073948, 47.27648]]
loader.close()
assert got[0] == expected[0]
assert is_equal_np(got[1], expected[1])
# grouped
dir_paths = ["./data/test/nifti/grouped/test"]
name = "images"
loader = NiftiFileLoader(dir_paths=dir_paths, name=name, grouped=True)
index = (0, 1)
array = loader.get_data(index)
got = [
np.shape(array),
[np.amax(array),
np.amin(array),
np.mean(array),
np.std(array)],
]
expected = [(64, 64, 60), [255.0, 0.0, 85.67942, 49.193127]]
loader.close()
assert got[0] == expected[0]
assert is_equal_np(got[1], expected[1])
# multi dirs
dir_paths = [
"./data/test/nifti/grouped/train", "./data/test/nifti/grouped/test"
]
name = "images"
loader = NiftiFileLoader(dir_paths=dir_paths, name=name, grouped=True)
index = (0, 1)
array = loader.get_data(index)
got = [
np.shape(array),
[np.amax(array),
np.amin(array),
np.mean(array),
np.std(array)],
]
expected = [(64, 64, 60), [255.0, 0.0, 85.67942, 49.193127]]
loader.close()
assert got[0] == expected[0]
assert is_equal_np(got[1], expected[1])
