Как построить валидацию изображений с помощью ImageDataGenerator?

Question

Я пытаюсь построить модель для классификации изображений, но не могу понять, как построить проверочные изображения с предсказанным классом (и вероятностью), как в этом руководстве:

https://www.tensorflow.org/tutorials/images/hub_with_keras#check_the_predictions

Поскольку я использую ImageDataGenerator, я могу получить информацию о предсказанных классах и вероятности, но я не могу получить сами изображения.

plt.figure(figsize=(10,9))
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(labels_batch[n])
  plt.axis('off')
_ = plt.suptitle("Model predictions")

Я получаю ValueError: не удалось передать входной массив из фигуры (400,80,80, 3) в форму (400)

Я знаю, что проблема в -> image_batch [n]

Пожалуйста, посоветуйте мне, как передать информацию об изображении в imshow ().

PS код для генератора изображений, который я использую:

img_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1/255)

image_batch = img_generator.flow_from_directory(
    path,
    target_size=(size, size),
    color_mode='rgb',
    class_mode='categorical',
    batch_size=batch_size
)