 b/test/test_generator.py
+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))