import os import random import shutil import numpy as np import pandas as pd import pytest from PIL import Image from keras_preprocessing.image import dataframe_iterator from keras_preprocessing.image import image_data_generator @pytest.fixture(scope='module') def all_test_images(): img_w = img_h = 20 rgb_images = [] rgba_images = [] gray_images = [] for n in range(8): bias = np.random.rand(img_w, img_h, 1) * 64 variance = np.random.rand(img_w, img_h, 1) * (255 - 64) imarray = np.random.rand(img_w, img_h, 3) * variance + bias im = Image.fromarray(imarray.astype('uint8')).convert('RGB') rgb_images.append(im) imarray = np.random.rand(img_w, img_h, 4) * variance + bias im = Image.fromarray(imarray.astype('uint8')).convert('RGBA') rgba_images.append(im) imarray = np.random.rand(img_w, img_h, 1) * variance + bias im = Image.fromarray( imarray.astype('uint8').squeeze()).convert('L') gray_images.append(im) return [rgb_images, rgba_images, gray_images] def test_dataframe_iterator(all_test_images, tmpdir): num_classes = 2 # save the images in the tmpdir count = 0 filenames = [] filepaths = [] filenames_without = [] for test_images in all_test_images: for im in test_images: filename = "image-{}.png".format(count) filename_without = "image-{}".format(count) filenames.append(filename) filepaths.append(os.path.join(str(tmpdir), filename)) filenames_without.append(filename_without) im.save(str(tmpdir / filename)) count += 1 df = pd.DataFrame({ "filename": filenames, "class": [str(random.randint(0, 1)) for _ in filenames], "filepaths": filepaths }) # create iterator iterator = dataframe_iterator.DataFrameIterator(df, str(tmpdir)) batch = next(iterator) assert len(batch) == 2 assert isinstance(batch[0], np.ndarray) assert isinstance(batch[1], np.ndarray) generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe(df, x_col='filepaths') df_iterator_dir = generator.flow_from_dataframe(df, str(tmpdir)) df_sparse_iterator = generator.flow_from_dataframe(df, str(tmpdir), class_mode="sparse") assert not np.isnan(df_sparse_iterator.classes).any() # check number of classes and images assert len(df_iterator.class_indices) == num_classes assert len(df_iterator.classes) == count assert set(df_iterator.filenames) == set(filepaths) assert len(df_iterator_dir.class_indices) == num_classes assert len(df_iterator_dir.classes) == count assert set(df_iterator_dir.filenames) == set(filenames) # test without shuffle _, batch_y = next(generator.flow_from_dataframe(df, str(tmpdir), shuffle=False, class_mode="sparse")) assert (batch_y == df['class'].astype('float')[:len(batch_y)]).all() # Test invalid use cases with pytest.raises(ValueError): generator.flow_from_dataframe(df, str(tmpdir), color_mode='cmyk') with pytest.raises(ValueError): generator.flow_from_dataframe(df, str(tmpdir), class_mode='output') with pytest.warns(DeprecationWarning): generator.flow_from_dataframe(df, str(tmpdir), has_ext=True) with pytest.warns(DeprecationWarning): generator.flow_from_dataframe(df, str(tmpdir), has_ext=False) def preprocessing_function(x): """This will fail if not provided by a Numpy array. Note: This is made to enforce backward compatibility. """ assert x.shape == (26, 26, 3) assert type(x) is np.ndarray return np.zeros_like(x) # Test usage as Sequence generator = image_data_generator.ImageDataGenerator( preprocessing_function=preprocessing_function) dir_seq = generator.flow_from_dataframe(df, str(tmpdir), target_size=(26, 26), color_mode='rgb', batch_size=3, class_mode='categorical') assert len(dir_seq) == np.ceil(count / 3) x1, y1 = dir_seq[1] assert x1.shape == (3, 26, 26, 3) assert y1.shape == (3, num_classes) x1, y1 = dir_seq[5] assert (x1 == 0).all() with pytest.raises(ValueError): x1, y1 = dir_seq[9] def test_dataframe_iterator_validate_filenames(all_test_images, tmpdir): # save the images in the paths count = 0 filenames = [] for test_images in all_test_images: for im in test_images: filename = 'image-{}.png'.format(count) im.save(str(tmpdir / filename)) filenames.append(filename) count += 1 df = pd.DataFrame({"filename": filenames + ['test.jpp', 'test.jpg']}) generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe(df, str(tmpdir), class_mode="input") assert len(df_iterator.filenames) == len(df['filename']) - 2 df_iterator = generator.flow_from_dataframe(df, str(tmpdir), class_mode="input", validate_filenames=False) assert len(df_iterator.filenames) == len(df['filename']) def test_dataframe_iterator_sample_weights(all_test_images, tmpdir): # save the images in the paths count = 0 filenames = [] for test_images in all_test_images: for im in test_images: filename = 'image-{}.png'.format(count) im.save(str(tmpdir / filename)) filenames.append(filename) count += 1 df = pd.DataFrame({"filename": filenames}) df['weight'] = ([2, 5] * len(df))[:len(df)] generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe(df, str(tmpdir), x_col="filename", y_col=None, shuffle=False, batch_size=5, weight_col='weight', class_mode="input") batch = next(df_iterator) assert len(batch) == 3 # (x, y, weights) # check if input and output have the same shape and they're the same assert(batch[0].all() == batch[1].all()) # check if the input and output images are not the same numpy array input_img = batch[0][0] output_img = batch[1][0] output_img[0][0][0] += 1 assert input_img[0][0][0] != output_img[0][0][0] assert np.array_equal(np.array([2, 5, 2, 5, 2]), batch[2]) # fail df['weight'] = (['2', '5'] * len(df))[:len(df)] with pytest.raises(TypeError): image_data_generator.ImageDataGenerator().flow_from_dataframe( df, weight_col='weight', class_mode="input" ) def test_dataframe_iterator_class_mode_input(all_test_images, tmpdir): # save the images in the paths count = 0 filenames = [] for test_images in all_test_images: for im in test_images: filename = 'image-{}.png'.format(count) im.save(str(tmpdir / filename)) filenames.append(filename) count += 1 df = pd.DataFrame({"filename": filenames}) generator = image_data_generator.ImageDataGenerator() df_autoencoder_iterator = generator.flow_from_dataframe(df, str(tmpdir), x_col="filename", y_col=None, class_mode="input") batch = next(df_autoencoder_iterator) # check if input and output have the same shape and they're the same assert np.allclose(batch[0], batch[1]) # check if the input and output images are not the same numpy array input_img = batch[0][0] output_img = batch[1][0] output_img[0][0][0] += 1 assert(input_img[0][0][0] != output_img[0][0][0]) df_autoencoder_iterator = generator.flow_from_dataframe(df, str(tmpdir), x_col="filename", y_col="class", class_mode="input") batch = next(df_autoencoder_iterator) # check if input and output have the same shape and they're the same assert(batch[0].all() == batch[1].all()) # check if the input and output images are not the same numpy array input_img = batch[0][0] output_img = batch[1][0] output_img[0][0][0] += 1 assert(input_img[0][0][0] != output_img[0][0][0]) def test_dataframe_iterator_class_mode_categorical_multi_label(all_test_images, tmpdir): # save the images in the paths filenames = [] count = 0 for test_images in all_test_images: for im in test_images: filename = 'image-{}.png'.format(count) im.save(str(tmpdir / filename)) filenames.append(filename) count += 1 label_opt = ['a', 'b', ['a'], ['b'], ['a', 'b'], ['b', 'a']] df = pd.DataFrame({ "filename": filenames, "class": [random.choice(label_opt) for _ in filenames[:-2]] + ['b', 'a'] }) generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe(df, str(tmpdir)) batch_x, batch_y = next(df_iterator) assert isinstance(batch_x, np.ndarray) assert len(batch_x.shape) == 4 assert isinstance(batch_y, np.ndarray) assert batch_y.shape == (len(batch_x), 2) for labels in batch_y: assert all(l in {0, 1} for l in labels) # on first 3 batches df = pd.DataFrame({ "filename": filenames, "class": [['b', 'a']] + ['b'] + [['c']] + [random.choice(label_opt) for _ in filenames[:-3]] }) generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe(df, str(tmpdir), shuffle=False) batch_x, batch_y = next(df_iterator) assert isinstance(batch_x, np.ndarray) assert len(batch_x.shape) == 4 assert isinstance(batch_y, np.ndarray) assert batch_y.shape == (len(batch_x), 3) for labels in batch_y: assert all(l in {0, 1} for l in labels) assert (batch_y[0] == np.array([1, 1, 0])).all() assert (batch_y[1] == np.array([0, 1, 0])).all() assert (batch_y[2] == np.array([0, 0, 1])).all() def test_dataframe_iterator_class_mode_multi_output(all_test_images, tmpdir): # save the images in the paths filenames = [] count = 0 for test_images in all_test_images: for im in test_images: filename = 'image-{}.png'.format(count) im.save(str(tmpdir / filename)) filenames.append(filename) count += 1 # fit both outputs are a single number df = pd.DataFrame({"filename": filenames}).assign( output_0=np.random.uniform(size=len(filenames)), output_1=np.random.uniform(size=len(filenames)) ) df_iterator = image_data_generator.ImageDataGenerator().flow_from_dataframe( df, y_col=['output_0', 'output_1'], directory=str(tmpdir), batch_size=3, shuffle=False, class_mode='multi_output' ) batch_x, batch_y = next(df_iterator) assert isinstance(batch_x, np.ndarray) assert len(batch_x.shape) == 4 assert isinstance(batch_y, list) assert len(batch_y) == 2 assert np.array_equal(batch_y[0], np.array(df['output_0'].tolist()[:3])) assert np.array_equal(batch_y[1], np.array(df['output_1'].tolist()[:3])) # if one of the outputs is a 1D array df['output_1'] = [np.random.uniform(size=(2, 2, 1)).flatten() for _ in range(len(df))] df_iterator = image_data_generator.ImageDataGenerator().flow_from_dataframe( df, y_col=['output_0', 'output_1'], directory=str(tmpdir), batch_size=3, shuffle=False, class_mode='multi_output' ) batch_x, batch_y = next(df_iterator) assert isinstance(batch_x, np.ndarray) assert len(batch_x.shape) == 4 assert isinstance(batch_y, list) assert len(batch_y) == 2 assert np.array_equal(batch_y[0], np.array(df['output_0'].tolist()[:3])) assert np.array_equal(batch_y[1], np.array(df['output_1'].tolist()[:3])) # if one of the outputs is a 2D array df['output_1'] = [np.random.uniform(size=(2, 2, 1)) for _ in range(len(df))] df_iterator = image_data_generator.ImageDataGenerator().flow_from_dataframe( df, y_col=['output_0', 'output_1'], directory=str(tmpdir), batch_size=3, shuffle=False, class_mode='multi_output' ) batch_x, batch_y = next(df_iterator) assert isinstance(batch_x, np.ndarray) assert len(batch_x.shape) == 4 assert isinstance(batch_y, list) assert len(batch_y) == 2 assert np.array_equal(batch_y[0], np.array(df['output_0'].tolist()[:3])) assert np.array_equal(batch_y[1], np.array(df['output_1'].tolist()[:3])) # fail if single column with pytest.raises(TypeError): image_data_generator.ImageDataGenerator().flow_from_dataframe( df, y_col='output_0', directory=str(tmpdir), class_mode='multi_output' ) def test_dataframe_iterator_class_mode_raw(all_test_images, tmpdir): # save the images in the paths filenames = [] count = 0 for test_images in all_test_images: for im in test_images: filename = 'image-{}.png'.format(count) im.save(str(tmpdir / filename)) filenames.append(filename) count += 1 # case for 1D output df = pd.DataFrame({"filename": filenames}).assign( output_0=np.random.uniform(size=len(filenames)), output_1=np.random.uniform(size=len(filenames)) ) df_iterator = image_data_generator.ImageDataGenerator().flow_from_dataframe( df, y_col='output_0', directory=str(tmpdir), batch_size=3, shuffle=False, class_mode='raw' ) batch_x, batch_y = next(df_iterator) assert isinstance(batch_x, np.ndarray) assert len(batch_x.shape) == 4 assert isinstance(batch_y, np.ndarray) assert batch_y.shape == (3,) assert np.array_equal(batch_y, df['output_0'].values[:3]) # case with a 2D output df_iterator = image_data_generator.ImageDataGenerator().flow_from_dataframe( df, y_col=['output_0', 'output_1'], directory=str(tmpdir), batch_size=3, shuffle=False, class_mode='raw' ) batch_x, batch_y = next(df_iterator) assert isinstance(batch_x, np.ndarray) assert len(batch_x.shape) == 4 assert isinstance(batch_y, np.ndarray) assert batch_y.shape == (3, 2) assert np.array_equal(batch_y, df[['output_0', 'output_1']].values[:3]) @pytest.mark.parametrize('validation_split,num_training', [ (0.25, 18), (0.50, 12), (0.75, 6), ]) def test_dataframe_iterator_with_validation_split(all_test_images, validation_split, num_training, tmpdir): num_classes = 2 # save the images in the tmpdir count = 0 filenames = [] filenames_without = [] for test_images in all_test_images: for im in test_images: filename = "image-{}.png".format(count) filename_without = "image-{}".format(count) filenames.append(filename) filenames_without.append(filename_without) im.save(str(tmpdir / filename)) count += 1 df = pd.DataFrame({"filename": filenames, "class": [str(random.randint(0, 1)) for _ in filenames]}) # create iterator generator = image_data_generator.ImageDataGenerator( validation_split=validation_split ) df_sparse_iterator = generator.flow_from_dataframe(df, str(tmpdir), class_mode="sparse") if np.isnan(next(df_sparse_iterator)[:][1]).any(): raise ValueError('Invalid values.') with pytest.raises(ValueError): generator.flow_from_dataframe( df, tmpdir, subset='foo') train_iterator = generator.flow_from_dataframe(df, str(tmpdir), subset='training') assert train_iterator.samples == num_training valid_iterator = generator.flow_from_dataframe(df, str(tmpdir), subset='validation') assert valid_iterator.samples == count - num_training # check number of classes and images assert len(train_iterator.class_indices) == num_classes assert len(train_iterator.classes) == num_training assert len(set(train_iterator.filenames) & set(filenames)) == num_training def test_dataframe_iterator_with_custom_indexed_dataframe(all_test_images, tmpdir): num_classes = 2 # save the images in the tmpdir count = 0 filenames = [] for test_images in all_test_images: for im in test_images: filename = "image-{}.png".format(count) filenames.append(filename) im.save(str(tmpdir / filename)) count += 1 # create dataframes classes = np.random.randint(num_classes, size=len(filenames)) classes = [str(c) for c in classes] df = pd.DataFrame({"filename": filenames, "class": classes}) df2 = pd.DataFrame({"filename": filenames, "class": classes}, index=np.arange(1, len(filenames) + 1)) df3 = pd.DataFrame({"filename": filenames, "class": classes}, index=filenames) # create iterators seed = 1 generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe( df, str(tmpdir), seed=seed) df2_iterator = generator.flow_from_dataframe( df2, str(tmpdir), seed=seed) df3_iterator = generator.flow_from_dataframe( df3, str(tmpdir), seed=seed) # Test all iterators return same pairs of arrays for _ in range(len(filenames)): a1, c1 = next(df_iterator) a2, c2 = next(df2_iterator) a3, c3 = next(df3_iterator) assert np.array_equal(a1, a2) assert np.array_equal(a1, a3) assert np.array_equal(c1, c2) assert np.array_equal(c1, c3) def test_dataframe_iterator_n(all_test_images, tmpdir): # save the images in the tmpdir count = 0 filenames = [] for test_images in all_test_images: for im in test_images: filename = "image-{}.png".format(count) filenames.append(filename) im.save(str(tmpdir / filename)) count += 1 # exclude first two items n_files = len(filenames) input_filenames = filenames[2:] # create dataframes classes = np.random.randint(2, size=len(input_filenames)) classes = [str(c) for c in classes] df = pd.DataFrame({"filename": input_filenames}) df2 = pd.DataFrame({"filename": input_filenames, "class": classes}) # create iterators generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe( df, str(tmpdir), class_mode=None) df2_iterator = generator.flow_from_dataframe( df2, str(tmpdir), class_mode='binary') # Test the number of items in iterators assert df_iterator.n == n_files - 2 assert df2_iterator.n == n_files - 2 def test_dataframe_iterator_absolute_path(all_test_images, tmpdir): # save the images in the tmpdir count = 0 file_paths = [] for test_images in all_test_images: for im in test_images: filename = "image-{:0>5}.png".format(count) file_path = str(tmpdir / filename) file_paths.append(file_path) im.save(file_path) count += 1 # prepare an image with a forbidden extension. file_path_fbd = str(tmpdir / 'image-forbid.fbd') shutil.copy(file_path, file_path_fbd) # create dataframes classes = np.random.randint(2, size=len(file_paths)) classes = [str(c) for c in classes] df = pd.DataFrame({"filename": file_paths}) df2 = pd.DataFrame({"filename": file_paths, "class": classes}) df3 = pd.DataFrame({"filename": ['image-not-exist.png'] + file_paths}) df4 = pd.DataFrame({"filename": file_paths + [file_path_fbd]}) # create iterators generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe( df, None, class_mode=None, shuffle=False, batch_size=1) df2_iterator = generator.flow_from_dataframe( df2, None, class_mode='binary', shuffle=False, batch_size=1) df3_iterator = generator.flow_from_dataframe( df3, None, class_mode=None, shuffle=False, batch_size=1) df4_iterator = generator.flow_from_dataframe( df4, None, class_mode=None, shuffle=False, batch_size=1) validation_split = 0.2 generator_split = image_data_generator.ImageDataGenerator( validation_split=validation_split ) df_train_iterator = generator_split.flow_from_dataframe( df, None, class_mode=None, shuffle=False, subset='training', batch_size=1) df_val_iterator = generator_split.flow_from_dataframe( df, None, class_mode=None, shuffle=False, subset='validation', batch_size=1) # Test the number of items in iterators assert df_iterator.n == len(file_paths) assert df2_iterator.n == len(file_paths) assert df3_iterator.n == len(file_paths) assert df4_iterator.n == len(file_paths) assert df_val_iterator.n == int(validation_split * len(file_paths)) assert df_train_iterator.n == len(file_paths) - df_val_iterator.n # Test flow_from_dataframe for i in range(len(file_paths)): a1 = next(df_iterator) a2, _ = next(df2_iterator) a3 = next(df3_iterator) a4 = next(df4_iterator) if i < df_val_iterator.n: a5 = next(df_val_iterator) else: a5 = next(df_train_iterator) assert np.array_equal(a1, a2) assert np.array_equal(a1, a3) assert np.array_equal(a1, a4) assert np.array_equal(a1, a5) def test_dataframe_iterator_with_subdirs(all_test_images, tmpdir): num_classes = 2 # create folders and subfolders paths = [] for cl in range(num_classes): class_directory = 'class-{}'.format(cl) classpaths = [ class_directory, os.path.join(class_directory, 'subfolder-1'), os.path.join(class_directory, 'subfolder-2'), os.path.join(class_directory, 'subfolder-1', 'sub-subfolder') ] for path in classpaths: tmpdir.join(path).mkdir() paths.append(classpaths) # save the images in the paths count = 0 filenames = [] for test_images in all_test_images: for im in test_images: # rotate image class im_class = count % num_classes # rotate subfolders classpaths = paths[im_class] filename = os.path.join( classpaths[count % len(classpaths)], 'image-{}.png'.format(count)) filenames.append(filename) im.save(str(tmpdir / filename)) count += 1 # create dataframe classes = np.random.randint(num_classes, size=len(filenames)) classes = [str(c) for c in classes] df = pd.DataFrame({"filename": filenames, "class": classes}) # create iterator generator = image_data_generator.ImageDataGenerator() df_iterator = generator.flow_from_dataframe( df, str(tmpdir), class_mode='binary') # Test the number of items in iterator assert df_iterator.n == len(filenames) assert set(df_iterator.filenames) == set(filenames) if __name__ == '__main__': pytest.main([__file__])