tensorflow_experiment_buid_data_generator.py

"""
Based on https://www.tensorflow.org/tutorials/images/cnn
"""

import pathlib
import tensorflow as tf

# Definition of the constant
BATCH_SIZE = 30
EPOCHS = 5
IMG_HEIGHT = 150
IMG_WIDTH = 150

# Definition of the folder
data_dir_train = ''# folder of the training file
data_dir_train = pathlib.Path(data_dir_train)

image_count_train = len(list(data_dir_train.glob('*/*.jpg')))

# Define the generator
train_image_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255) # Generator for our training data

# Build the generator
train_data_gen = train_image_generator.flow_from_directory(batch_size=BATCH_SIZE,
                                                       directory=data_dir_train,
                                                       shuffle=True,
                                                       target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                       class_mode='binary')