def test_should_produce_augmented_samples_given_batch_size(self):
images = np.random.rand(20, 64, 64, 1)
labels = np.random.rand(20)
generator = DataGenerator().fit(images, labels)
batch_size = 10
batch, _ = generator.get_next_batch(batch_size)
self.assertEqual(batch.shape[0], batch_size)
self.assertEqual(images.shape[1:], batch.shape[1:])
def test_should_resize_images_to_given_target_dimension(self):
images = np.random.rand(20, 64, 64, 3)
labels = np.random.rand(20)
generator = DataGenerator().fit(images, labels)
batch, _ = generator.get_next_batch(10, target_dimensions=(28, 28))
self.assertEqual(batch.shape, (10, 28, 28, 3))
