from pathlib import Path
import numpy as np
import tensorflow as tf
from tensorflow import keras

# Load the trained model
model = keras.models.load_model('Network/trained_model.h5')

# Load the class names
class_names = ['table', 'done', 'order']

# Load and preprocess the validation dataset
data_dir = "Network/Training/"
image_size = (100, 100)
batch_size = 32

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=image_size,
    batch_size=batch_size,
)

# Select 20 random images from the validation set
val_images = []
val_labels = []
for images, labels in val_ds.unbatch().shuffle(1000).take(60):
    val_images.append(images)
    val_labels.append(labels)

# Make predictions on the random images
errorcount = 0
for i in range(60):
    test_image = val_images[i]
    test_label = val_labels[i]
    test_image = np.expand_dims(test_image, axis=0)
    test_image = test_image / 255.0  # Normalize the image

    # Make predictions
    predictions = model.predict(test_image)
    predicted_class_index = np.argmax(predictions[0])
    predicted_class = class_names[predicted_class_index]

    true_class = class_names[test_label]

    direct = 'Network/Results/'
    filename = predicted_class + str(i) + '.jpeg'
    Path(direct).mkdir(parents=True, exist_ok=True)
    tf.keras.preprocessing.image.save_img(direct+filename, val_images[i])
    if predicted_class != true_class:
        errorcount += 1
    print('Image', i+1)
    print('True class:', true_class)
    print('Predicted class:', predicted_class)
    print()
print('Error count: ', errorcount)