import os
from unittest import TestCase
import numpy as np
from fetal_net.data import add_data_to_storage, create_data_file
from fetal_net.generator import get_multi_class_labels, get_training_and_validation_generators
from fetal_net.augment import generate_permutation_keys, permute_data, reverse_permute_data
class TestDataGenerator(TestCase):
def setUp(self):
self.tmp_files = list()
self.data_file = None
def tearDown(self):
if self.data_file:
self.data_file.close()
self.rm_tmp_files()
def create_data_file(self, n_samples=20, len_x=5, len_y=5, len_z=10, n_channels=1):
self.data_file_path = "./temporary_data_test_file.h5"
self.training_keys_file = "./temporary_training_keys_file.pkl"
self.validation_keys_file = "./temporary_validation_keys_file.pkl"
self.tmp_files = [self.data_file_path, self.training_keys_file, self.validation_keys_file]
self.rm_tmp_files()
self.n_samples = n_samples
self.n_channels = n_channels
self.n_labels = 1
image_shape = (len_x, len_y, len_z)
data_size = self.n_samples * self.n_channels * len_x * len_y * len_z
data = np.asarray(np.arange(data_size).reshape((self.n_samples, self.n_channels, len_x, len_y, len_z)),
dtype=np.int16)
self.assertEqual(data.shape[-3:], image_shape)
truth = (data[:, 0] == 3).astype(np.int8).reshape(data.shape[0], 1, data.shape[2], data.shape[3], data.shape[4])
affine = np.diag(np.ones(4))
affine[:, -1] = 1
self.data_file, data_storage, truth_storage, affine_storage = create_data_file(self.data_file_path,
self.n_channels, self.n_samples,
image_shape)
for index in range(self.n_samples):
add_data_to_storage(data_storage, truth_storage, affine_storage,
np.concatenate([data[index], truth[index]], axis=0), affine=affine,
n_channels=self.n_channels,
truth_dtype=np.int16)
self.assertTrue(np.all(data_storage[index] == data[index]))
self.assertTrue(np.all(truth_storage[index] == truth[index]))
def rm_tmp_files(self):
for tmp_file in self.tmp_files:
if os.path.exists(tmp_file):
os.remove(tmp_file)
def test_multi_class_labels(self):
n_labels = 5
labels = np.arange(1, n_labels+1)
x_dim = 3
label_map = np.asarray([[[np.arange(n_labels+1)] * x_dim]])
binary_labels = get_multi_class_labels(label_map, n_labels, labels)
for label in labels:
self.assertTrue(np.all(binary_labels[:, label - 1][label_map[:, 0] == label] == 1))
def test_get_training_and_validation_generators(self):
self.create_data_file()
validation_split = 0.8
batch_size = 3
validation_batch_size = 3
generators = get_training_and_validation_generators(data_file=self.data_file, batch_size=batch_size,
n_labels=self.n_labels,
training_keys_file=self.training_keys_file,
validation_keys_file=self.validation_keys_file,
data_split=validation_split,
validation_batch_size=validation_batch_size,
skip_blank=False)
training_generator, validation_generator, n_training_steps, n_validation_steps = generators
self.verify_generator(training_generator, n_training_steps, batch_size,
np.round(validation_split * self.n_samples))
self.verify_generator(validation_generator, n_validation_steps, validation_batch_size,
np.round((1 - validation_split) * self.n_samples))
self.data_file.close()
self.rm_tmp_files()
def verify_generator(self, generator, steps, batch_size, expected_samples):
# check that the generator covers all the samples
n_validation_samples = 0
validation_samples = list()
for i in range(steps):
x, y = next(generator)
hash_x = hash(str(x))
self.assertNotIn(hash_x, validation_samples)
validation_samples.append(hash_x)
n_validation_samples += x.shape[0]
if i + 1 != steps:
self.assertEqual(x.shape[0], batch_size)
self.assertEqual(n_validation_samples, expected_samples)
def test_patch_generators(self):
self.create_data_file(len_x=4, len_y=4, len_z=4)
validation_split = 0.8
batch_size = 10
validation_batch_size = 3
patch_shape = (2, 2, 2)
generators = get_training_and_validation_generators(self.data_file, batch_size, self.n_labels,
self.training_keys_file, self.validation_keys_file,
patch_shape=patch_shape, data_split=validation_split,
validation_batch_size=validation_batch_size,
skip_blank=False)
training_generator, validation_generator, n_training_steps, n_validation_steps = generators
expected_training_samples = int(np.round(self.n_samples * validation_split)) * 2**3
self.verify_generator(training_generator, n_training_steps, batch_size, expected_training_samples)
expected_validation_samples = int(np.round(self.n_samples * (1 - validation_split))) * 2**3
self.verify_generator(validation_generator, n_validation_steps, validation_batch_size,
expected_validation_samples)
self.data_file.close()
self.rm_tmp_files()
def test_random_patch_start(self):
self.create_data_file(len_x=10, len_y=10, len_z=10)
validation_split = 0.8
batch_size = 10
validation_batch_size = 3
patch_shape = (5, 5, 5)
random_start = (3, 3, 3)
overlap = 2
generators = get_training_and_validation_generators(self.data_file, batch_size, self.n_labels,
self.training_keys_file, self.validation_keys_file,
patch_shape=patch_shape, data_split=validation_split,
validation_batch_size=validation_batch_size,
skip_blank=False)
training_generator, validation_generator, n_training_steps, n_validation_steps = generators
expected_training_samples = int(np.round(self.n_samples * validation_split)) * 2**3
self.verify_generator(training_generator, n_training_steps, batch_size, expected_training_samples)
expected_validation_samples = int(np.round(self.n_samples * (1 - validation_split))) * 4**3
self.verify_generator(validation_generator, n_validation_steps, validation_batch_size,
expected_validation_samples)
self.data_file.close()
self.rm_tmp_files()
def test_unique_permutations(self):
permutations = list()
shape = (2, 3, 3, 3)
data = np.arange(54).reshape(shape)
for key in generate_permutation_keys():
permutations.append(permute_data(data, key))
for array in permutations[:-1]:
self.assertTrue(permutations[-1].shape == shape)
self.assertFalse(np.all(array == permutations[-1]))
self.assertEqual(np.sum(data), np.sum(permutations[-1]))
def test_n_permutations(self):
self.assertEqual(len(generate_permutation_keys()), 48)
def test_generator_with_permutations(self):
self.create_data_file(len_x=5, len_y=5, len_z=5, n_channels=5)
batch_size = 2
generators = get_training_and_validation_generators(self.data_file, batch_size, self.n_labels,
self.training_keys_file, self.validation_keys_file)
training_generator, validation_generator, n_training_steps, n_validation_steps = generators
for x in training_generator:
break
self.rm_tmp_files()
def test_reverse_permutation(self):
data_shape = (4, 32, 32, 32)
data = np.arange(np.prod(data_shape)).reshape(data_shape)
for permutation_key in generate_permutation_keys():
permuted_data = permute_data(data, permutation_key)
reversed_permutation = reverse_permute_data(permuted_data, permutation_key)
self.assertTrue(np.all(data == reversed_permutation))
