211 lines
6.6 KiB
Python
211 lines
6.6 KiB
Python
from albumentations.augmentations.utils import P
|
|
from PythonDataTransferMQRPC import PythonDataTransferMQRPC
|
|
import albumentations as A
|
|
import EgLib as El
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
from sklearn.model_selection import train_test_split
|
|
from PIL import Image
|
|
from tensorflow.keras.utils import Sequence
|
|
import time
|
|
|
|
class DataGenerator(Sequence):
|
|
def __init__(self, x_set, y_set, batch_size):
|
|
self.x, self.y = x_set, y_set
|
|
self.batch_size = batch_size
|
|
|
|
def __len__(self):
|
|
return int(np.ceil(len(self.x) / float(self.batch_size)))
|
|
|
|
def __getitem__(self, idx):
|
|
batch_x = self.x[idx * self.batch_size:(idx + 1) * self.batch_size]
|
|
batch_y = self.y[idx * self.batch_size:(idx + 1) * self.batch_size]
|
|
return batch_x, batch_y
|
|
|
|
transform = A.Compose([
|
|
|
|
A.HorizontalFlip(p=0.5),
|
|
A.VerticalFlip(p=0.5),
|
|
A.Rotate(p=0.8),
|
|
|
|
])
|
|
classes = 25
|
|
|
|
models = {}
|
|
|
|
|
|
|
|
class HandlerClass:
|
|
|
|
def __init__(self):
|
|
self.model = None
|
|
self.active_model = None
|
|
self.val_dataset = None
|
|
|
|
self.train_dataset = None
|
|
|
|
self.Y_test = None
|
|
self.Y_train = None
|
|
self.X_train = None
|
|
self.X_test = None
|
|
self.const={}
|
|
|
|
def activate_model(self, name):
|
|
|
|
self.model = models[name]
|
|
self.active_model = name
|
|
|
|
|
|
|
|
return 0
|
|
|
|
def start_learning(self, path):
|
|
print("start learning")
|
|
target_size = self.get_accept_size()
|
|
images, labels = El.load_hawkeye(path, target_size=(target_size[1], target_size[2]))
|
|
|
|
images_augmented = []
|
|
labels_augmented = []
|
|
|
|
for _ in range(15):
|
|
for i, l in zip(images, labels):
|
|
transformed = transform(image=i, mask=l)
|
|
images_augmented.append(transformed["image"])
|
|
labels_augmented.append(transformed["mask"])
|
|
images, labels = np.array(images_augmented), np.array(labels_augmented)
|
|
labels = np.squeeze(labels[:, :, :, 0])
|
|
labels = tf.one_hot(labels, classes)
|
|
|
|
images_normalized = images / 255.
|
|
self.X_train, self.X_test, self.Y_train, self.Y_test = train_test_split(images_normalized, labels.numpy(),
|
|
test_size=0.33, random_state=42)
|
|
|
|
self.train_dataset = DataGenerator(self.X_train, self.Y_train, 16)
|
|
self.val_dataset = DataGenerator(self.X_test, self.Y_test, 16)
|
|
|
|
return 0
|
|
|
|
def create_model(self, name, additional_layers, classes):
|
|
input_layer = layers.Input(shape=(128,128,3))
|
|
x = layers.Conv2D(16, 7,strides=1, padding='same',activation="relu")(input_layer)
|
|
x = layers.MaxPooling2D()(x)
|
|
x = residual_block(x,16,5)
|
|
x = layers.MaxPooling2D()(x)
|
|
|
|
for _ in range(additional_layers):
|
|
x = residual_block(x,16,3)
|
|
|
|
x = residual_block(x,16,3)
|
|
x = layers.UpSampling2D()(x)
|
|
x = residual_block(x,16,3)
|
|
x = layers.UpSampling2D()(x)
|
|
x = layers.Conv2D(9, 3, padding='same',activation="softmax")(x)
|
|
|
|
|
|
model = Model(inputs=input_layer,outputs=x)
|
|
model.summary()
|
|
|
|
model.compile(loss="categorical_crossentropy",optimizer=tf.optimizers.Adam(learning_rate=0.01),metrics=["accuracy"])
|
|
|
|
def save_model(self, path):
|
|
self.model.save(path)
|
|
return 0
|
|
|
|
def load_model(self, path, name):
|
|
models[name]=tf.keras.models.load_model(path)
|
|
|
|
|
|
self.const[name]=None
|
|
|
|
return 0
|
|
|
|
def get_accept_size(self):
|
|
return list(self.model.input_shape)
|
|
|
|
def get_output_size(self):
|
|
return list(self.model.layers[-1].output.shape[1:4])
|
|
|
|
def predict(self, data):
|
|
target_size = self.get_accept_size()
|
|
image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.
|
|
if self.const[self.active_model] is not None:
|
|
predictions = self.model.predict([np.array([image]),self.const])
|
|
else:
|
|
predictions = self.model.predict(np.array([image]))
|
|
|
|
|
|
|
|
if predictions[0].ndim==2:
|
|
classed = predictions[0]
|
|
else:
|
|
classed = tf.argmax(predictions[0], axis=2)
|
|
|
|
print(classed.shape)
|
|
|
|
return classed.numpy().reshape(-1).tolist()
|
|
|
|
def predict_anomaly(self,data):
|
|
target_size = self.get_accept_size()
|
|
image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.
|
|
inputs=[np.array([image]),np.array([self.const[self.active_model]])]
|
|
|
|
predictions = self.model.predict(inputs)
|
|
print(predictions[0][0][0].shape)
|
|
print(predictions[1].shape)
|
|
|
|
image=predictions[0][0][0]
|
|
|
|
|
|
normalized_image=(image-image.min())/(image.max()-image.min())
|
|
defect_probability=predictions[1][0][1]*100
|
|
squized_image=(normalized_image*255).reshape(-1)
|
|
|
|
squized_prop=np.array([defect_probability*100]).astype(int)
|
|
|
|
res = np.concatenate((squized_prop,(normalized_image*255).astype(int).reshape(-1))).astype(np.uint8).tolist()
|
|
print(len(res))
|
|
return res
|
|
|
|
def predict_raw(self, data):
|
|
target_size = self.get_accept_size()
|
|
image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.
|
|
|
|
# measure time
|
|
strart_time=time.time()
|
|
|
|
if self.const[self.active_model] is not None:
|
|
predictions = self.model.predict([np.array([image]),self.const])
|
|
else:
|
|
print("test")
|
|
predictions = self.model(np.array([image])).numpy()
|
|
|
|
print(time.time()-strart_time)
|
|
|
|
if predictions[0].ndim==2:
|
|
classed = predictions[0][:,:,np.newaxis].clip(0,255)
|
|
return (classed).astype(np.uint8).reshape(-1).tolist()
|
|
else:
|
|
classed = predictions[0][:,:,:]
|
|
|
|
print(classed.max())
|
|
print(classed.min())
|
|
return (classed*255).astype(np.uint8).reshape(-1).tolist()
|
|
|
|
|
|
|
|
def continue_learning(self):
|
|
print("continue learning")
|
|
schedule = tf.keras.callbacks.LearningRateScheduler(
|
|
lambda epoch: 0.001, verbose=0)
|
|
|
|
model_history = self.model.fit(self.train_dataset, epochs=1,
|
|
validation_data=self.val_dataset,
|
|
callbacks=[schedule])
|
|
tf.keras.backend.clear_session()
|
|
return model_history.history["val_loss"][0]
|
|
|
|
|
|
data_transfer = PythonDataTransferMQRPC()
|
|
|
|
data_transfer.wait_for_data_and_process(HandlerClass())
|