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())