import os
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']

# Path to the folder containing test images
test_images_folder = 'Network/Testing/'

# Iterate over the test images
i = 0
errorcount = 0
for folder_name in os.listdir(test_images_folder):
    folder_path = os.path.join(test_images_folder, folder_name)
    if os.path.isdir(folder_path):
        print('Testing images in folder:', folder_name)

        # True class based on folder name
        if folder_name == 'Empty':
            true_class = 'table'
        elif folder_name == 'Food':
            true_class = 'done'
        elif folder_name == 'People':
            true_class = 'order'

        # Iterate over the files in the subfolder
        for filename in os.listdir(folder_path):
            if filename.endswith('.jpg') or filename.endswith('.jpeg'):
                i += 1
                # Load and preprocess the test image
                image_path = os.path.join(folder_path, filename)
                test_image = keras.preprocessing.image.load_img(
                    image_path, target_size=(100, 100))
                test_image = keras.preprocessing.image.img_to_array(test_image)
                test_image = np.expand_dims(test_image, axis=0)
                test_image = test_image / 255.0  # Normalize the image

                # Reshape the image array to (1, height, width, channels)
                test_image = np.reshape(test_image, (1, 100, 100, 3))

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

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