admin/HawkeyeVision
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

WSL + API support

Browse Source
This commit is contained in:
meelstorm
2025-08-20 15:57:50 +02:00
parent 2f570545ab
commit af8ee449a0
7 changed files with 371 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

210
PythonModelAPI/Segment128Half.py Normal file
View File

@@ -0,0 +1,210 @@
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())