import tensorflow as tf
import glob
import numpy as np
import os

from PIL import Image
from sklearn.model_selection import train_test_split
from tensorflow import keras

images = []
labels = []

class_names = ['fragile', 'priority', 'skull']

image_dir = 'C:\\Users\\mateu\\Desktop\\SI projekt\\resources'
image_size = (32, 32)

for class_index, class_name in enumerate(class_names):
    class_path = os.path.join(image_dir, class_name)
    file_names = os.listdir(class_path)
    print(class_path)
    for file_name in file_names:
        image_path = os.path.join(class_path, file_name)
        image = Image.open(image_path).convert('L')
        image = image.resize(image_size)
        image_array = np.array(image) / 255.0
        images.append(image_array)
        labels.append(class_index)

images = np.array(images)
labels = np.array(labels)

X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)

num_classes = 3
y_train_encoded = keras.utils.to_categorical(y_train, num_classes)
y_test_encoded = keras.utils.to_categorical(y_test, num_classes)

model = keras.Sequential([
    keras.layers.Flatten(input_shape=image_size),
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(num_classes, activation='softmax')
])

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

batch_size = 32
epochs = 10

model.fit(X_train, y_train_encoded, batch_size=batch_size, epochs=epochs,
          validation_data=(X_test, y_test_encoded))

test_loss, test_acc = model.evaluate(X_test, y_test_encoded)
print('Test accuracy:', test_acc)

predictions = model.predict(X_test)

model.save('resources/model.h5')