import os
import numpy as np
import tensorflow as tf
from tensorflow import keras

train_data_dir = "Training/"

train_ds = tf.keras.utils.image_dataset_from_directory(train_data_dir, validation_split=0.2,
                                                       subset="training", seed=123, batch_size=32, image_size=(100, 100))
val_ds = tf.keras.utils.image_dataset_from_directory(train_data_dir, validation_split=0.2,
                                                     subset="validation", seed=123, batch_size=32, image_size=(100, 100))

model = keras.models.load_model("trained_model")
predictions = model.predict(val_ds.take(32))


classNames = ['Empty', 'Food','People']

# Make predictions
direct = ''
i = 0
for image, _ in val_ds.take(32):
    predicted_class_index = np.argmax(predictions[i])
    predicted_class = classNames[predicted_class_index]
    filename = predicted_class + str(i) + '.jpeg'
    tf.keras.preprocessing.image.save_img(direct+filename, image[0])
    print('Predicted class:', predicted_class)
    i += 1

#direct = ''

#i = 0
#for image, _ in val_ds.take(32):
#    predictedLabel = int(predictions[i] >= 0.5)
#
#    filename = classNames[predictedLabel] + str(i) + '.jpeg'
#    tf.keras.preprocessing.image.save_img(direct+filename, image[0])
#    i += 1