Tensorflow CNN Image Classification - Using ImageDataGenerator and then Next&Model.fit gives error

Question

I have a CNN model, which is basically processing images and classifying them at the end. There are four class labels, which are UN, D1, D2 and D3. If you look at train_batches, you will see that it already labels them as an integer from 0 to 3. Before I send those images into CNN model, I been doing preprocessing. Those batches give me correct number of images and classes. However, when I'd like to plot them to see if I am doing anything wrong (using next function) or when I run model.fit_generator(x=train_batches, validation_data =valid_batches, epochs=5, verbose=2), it says "IndexError: index 1 is out of bounds for axis 2 with size 1".

I was suspecting there might be some "number of class" incompatibility but I could not figure it out.

train_batches = ImageDataGenerator(preprocessing_function=tf.keras.applications.vgg16.preprocess_input).\
    flow_from_directory(directory=train_path, color_mode='grayscale', target_size=(24, 1000),
                        classes=['UN', 'D1', 'D2', 'D3'], shuffle=True, batch_size=5)

test_batches = ImageDataGenerator(preprocessing_function=tf.keras.applications.vgg16.preprocess_input).\
    flow_from_directory(directory=test_path, color_mode='grayscale', target_size=(24, 1000),
                        class_mode=None, shuffle=False, batch_size=1)

valid_batches = ImageDataGenerator(preprocessing_function=tf.keras.applications.vgg16.preprocess_input).\
    flow_from_directory(directory=valid_path, color_mode='grayscale', target_size=(24, 1000),
                        classes=['UN', 'D1', 'D2', 'D3'], shuffle=True, batch_size=5)

assert train_batches.n == 240
assert test_batches.n == 40
assert valid_batches.n == 41
assert train_batches.num_classes == valid_batches.num_classes == test_batches.num_classes == 4

model.compile(optimizer=Adam(learning_rate=0.0001), loss='SparseCategoricalCrossentropy', metrics=['accuracy'])
step_size_train = train_batches.n // train_batches.batch_size
step_size_valid = valid_batches.n // valid_batches.batch_size
model.fit_generator(generator=train_batches, steps_per_epoch=step_size_train, validation_data=valid_batches,
                    validation_steps=step_size_valid, epochs=10)

Answer 1

You are getting the error because tf.keras.applications.vgg16.preprocess_input takes an input tensor with 3 channels, according to its documentation. You don't need this function since you're training your model from scratch. Passing rescale=1/255 in the ImageDataGenerator call will be fine for basic preprocessing.

train_batches = ImageDataGenerator(rescale=1/255).flow_from_directory(directory=train_path,
                color_mode='grayscale', 
                target_size=(256,256),
                classes=['UN', 'D1', 'D2', 'D3'], 
                shuffle=True, 
                batch_size=5)

                

Also, as you are working with images, the input shape in the model will be the shape of the image. For example (256,256,3) and the target size will be (256,256)

Let us know if the issue still persists. Thanks!