Python, сохраняйте количество обучающих и проверочных изображений, полученных из ImageDataGenerator, в переменных для дальнейшего использования

Question

У меня есть следующий python код, который я запускаю в Блокнотах Google Colab, который использует ImageDataGenerator для разделения набора разработки на обучение и проверку.

datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale=1./255, 
    validation_split=0.35)

train_data_gen = datagen.flow_from_directory(
    base_dir,
    shuffle=True,
    target_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size,
    subset='training' 
    )

val_data_gen = datagen.flow_from_directory(
    base_dir,
    target_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size, 
    subset='validation')

Вывод кода приведен ниже:

Found 4708 images belonging to 40 classes.
Found 2512 images belonging to 40 classes.

Я хочу сохранить значения (4708 и 2512) в переменных (например, train_size и val_size) для дальнейшего использования. Может кто-нибудь объяснить, как я это сделаю?

Answer 1

Вы не можете получить эти числа как результат flow_from_directory().

Что вы можете сделать вместо этого (поскольку я предполагаю, что вы хотите ввести steps_per_epoch = training_set_length // batch_size и то же самое для validation_steps), вы можете написать некоторый произвольный код Python для решения этой проблемы.

Например, вы можете получить число по:

class_subfolders = [os.path.join(main_folder,class_subfolder) for class_subfolder in os.scandir(main_folder)]
#Either training or validation, depending on what 'main_folder' you feed
number_of_samples = sum([len(os.listdir(class_subfolder)) for class_subfolder in class_subfolders]