python keras как предсказать значение для всех элементов массива

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

test_dataset = image_dataset_from_directory('D:\\neyro_1\\prog\\new\\',
                                    batch_size=batch_size,
                                    image_size=image_size) 

test_dataset = test_dataset.prefetch(buffer_size=AUTOTUNE)

predictions = model.predict(test_dataset)

проблема состоит в том, что при предсказании результаты выводятся в случайном порядке, т.е. predictions[0][0] не соответствует 1 картинке в папке и так далее. Вопрос: каким образом мне определить, какое предсказание какой картинке соответствует?(по отдельности предсказывать для каждой картинки - не вариант) Если знаете как мне помочь, пожалуйста помогите!)