59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
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')
|