Traktor/app/neural_network.py
2021-06-02 14:06:02 +02:00

113 lines
4.0 KiB
Python

#!/usr/bin/python3
import os
from tensorflow.keras.models import Sequential, save_model, load_model
from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras.losses import sparse_categorical_crossentropy
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow import keras as k
import numpy as np
from app.base_field import BaseField
from config import *
class NeuralNetwork:
def __init__(self):
# Model config
self.batch_size = 25
self.img_width, self.img_height, self.img_num_channels = 25, 25, 3
self.loss_function = sparse_categorical_crossentropy
self.no_classes = 8
self.no_epochs = 40
self.optimizer = Adam()
self.verbosity = 1
# Determine shape of the data
self.input_shape = (self.img_width, self.img_height, self.img_num_channels)
# labels
self.labels = ["cabbage", "carrot", "corn", "lettuce", "paprika", "potato", "sunflower", "tomato"]
def init_model(self) -> None:
if not self.model_dir_is_empty():
# Load the model
self.model = load_model(
os.path.join(RESOURCE_DIR, "saved_model"),
custom_objects=None,
compile=True
)
else:
# Create the model
self.model = Sequential()
self.model.add(Conv2D(16, kernel_size=(5, 5), activation='relu', input_shape=self.input_shape))
self.model.add(Conv2D(32, kernel_size=(5, 5), activation='relu'))
self.model.add(Conv2D(64, kernel_size=(5, 5), activation='relu'))
self.model.add(Conv2D(128, kernel_size=(5, 5), activation='relu'))
self.model.add(Flatten())
self.model.add(Dense(16, activation='relu'))
self.model.add(Dense(self.no_classes, activation='softmax'))
self.model.compile(loss=self.loss_function,
optimizer=self.optimizer,
metrics=['accuracy'])
# Start training
self.model.fit(
self.train_datagen,
epochs=self.no_epochs,
shuffle=False)
# Display a model summary
# self.model.summary()
def load_images(self) -> None:
# Create a generator
self.train_datagen = ImageDataGenerator(
rescale=1. / 255
)
self.train_datagen = self.train_datagen.flow_from_directory(
TRAINING_SET_DIR,
save_to_dir=ADAPTED_IMG_DIR,
save_format='jpeg',
batch_size=self.batch_size,
target_size=(25, 25),
class_mode='sparse')
def predict(self, field: BaseField) -> str:
print(field.get_img_path())
# corn_img_path = os.path.join(RESOURCE_DIR,'corn.png')
loaded_image = k.preprocessing.image.load_img(field.get_img_path(),
target_size=(
self.img_width, self.img_height, self.img_num_channels))
# convert to array and resample dividing by 255
img_array = k.preprocessing.image.img_to_array(loaded_image) / 255.
# add sample dimension. the predictor is expecting (1, CHANNELS, IMG_WIDTH, IMG_HEIGHT)
img_np_array = np.expand_dims(img_array, axis=0)
# print(img_np_array)
predictions = self.model.predict(img_np_array)
prediction = np.argmax(predictions[0])
label = self.labels[prediction]
print(f'Ground truth: {type(field).__name__} - Prediction: {label}')
return label
def model_dir_is_empty(self) -> bool:
if len(os.listdir(MODEL_DIR)) == 0:
return True
return False
def check(self, field: BaseField) -> str:
self.load_images()
self.init_model()
prediction = self.predict(field)
# Saving model
if self.model_dir_is_empty():
save_model(self.model, MODEL_DIR)
return prediction