71 lines
2.9 KiB
Python
71 lines
2.9 KiB
Python
import cv2
|
|
import matplotlib.pyplot as plt
|
|
import os
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
|
|
import tensorflow as tf
|
|
from tensorflow.keras import models, layers
|
|
import numpy as np
|
|
|
|
BATCH_SIZE = 1 # 32
|
|
IMG_SIZE = (32, 32)
|
|
# MODEL_PATH = "C:/Users/annas/Desktop/projekty_studia/sztuczna/AIprojekt-wozek/neural_network/image_classifier.model"
|
|
MODEL_PATH = "neural_network/image_classifier.model"
|
|
|
|
|
|
class NeuralNetwork:
|
|
# def __init__(s):
|
|
# self.photo = photo
|
|
|
|
def load_data(self):
|
|
path_markings = "resources/package_markings_data/"
|
|
path_val_markings = "resources/package_markings_val/"
|
|
|
|
training_images = tf.keras.utils.image_dataset_from_directory(path_markings, shuffle=True,
|
|
batch_size=BATCH_SIZE,
|
|
image_size=IMG_SIZE)
|
|
|
|
validation_data = tf.keras.utils.image_dataset_from_directory(path_val_markings, shuffle=True,
|
|
batch_size=BATCH_SIZE,
|
|
image_size=IMG_SIZE)
|
|
AUTOTUNE = tf.data.AUTOTUNE
|
|
train_dataset = training_images.prefetch(buffer_size=AUTOTUNE)
|
|
validation_dataset = validation_data.prefetch(buffer_size=AUTOTUNE)
|
|
return train_dataset, validation_dataset
|
|
|
|
def train(self):
|
|
model = models.Sequential()
|
|
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)))
|
|
model.add(layers.MaxPooling2D((2, 2)))
|
|
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
|
|
model.add(layers.MaxPooling2D((2, 2)))
|
|
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
|
|
model.add(layers.Flatten())
|
|
model.add(layers.Dense(64, activation='relu'))
|
|
model.add(layers.Dense(2, activation='softmax'))
|
|
|
|
model.compile(optimizer='adam', loss="sparse_categorical_crossentropy", metrics=['accuracy'])
|
|
train_dataset, validation_dataset = self.load_data()
|
|
model.fit(train_dataset, epochs=10, validation_data=(validation_dataset))
|
|
|
|
loss, accuracy = model.evaluate(validation_dataset)
|
|
print(f"Loss = {loss}", f"Accuracy = {accuracy}")
|
|
model.save(MODEL_PATH)
|
|
|
|
def use_neural_net(self, image):
|
|
model = models.load_model(MODEL_PATH)
|
|
|
|
img = cv2.imread(image)
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
plt.imshow(img, cmap=plt.cm.binary)
|
|
plt.show()
|
|
class_names = ["fragile", "priority"]
|
|
prediction = model.predict(np.array([img]))
|
|
index = np.argmax(prediction)
|
|
prediction = class_names[index]
|
|
print(prediction)
|
|
return prediction
|
|
|
|
|
|
# net = NeuralNetwork("C:/Users/annas/Desktop/projekty_studia/sztuczna/AIprojekt-wozek/priority (2).png")
|
|
# net.train()
|
|
# print(net.use_neural_net()) |