diff --git a/Marcin.py b/Marcin.py index c62a460..db3fb80 100644 --- a/Marcin.py +++ b/Marcin.py @@ -1,5 +1,7 @@ -# import keras -import numpy +import numpy as np +import cv2 +import pygame + class main(): def __init__(self,traktor,field,ui,path): @@ -8,5 +10,62 @@ class main(): self.ui = ui self.path = path + def get_output_layers(self,net): + layer_names = net.getLayerNames() + output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()] + return output_layers + + def recognition(self,photo): + image = photo + + Width = image.shape[1] + Height = image.shape[0] + scale = 0.00392 + + with open("si.names", 'r') as f: + classes = [line.strip() for line in f.readlines()] + + COLORS = np.random.uniform(0, 255, size=(len(classes), 3)) + + net = cv2.dnn.readNet("si_final.weights", "si.cfg") + + blob = cv2.dnn.blobFromImage(image, scale, (416, 416), (0, 0, 0), True, crop=False) + + net.setInput(blob) + + outs = net.forward(self.get_output_layers(net)) + + class_ids = [] + confidences = [] + boxes = [] + conf_threshold = 0.5 + nms_threshold = 0.4 + + for out in outs: + for detection in out: + scores = detection[5:] + class_id = np.argmax(scores) + confidence = scores[class_id] + if confidence > 0.5: + class_ids.append(class_id) + print(class_id) + print(scores) + return class_id + def main(self): - pass + self.pole = self.ui.field_images[self.field.get_value(self.traktor.get_poz())] + self.img = pygame.surfarray.array3d(self.pole) + self.img = self.img.transpose([1,0,2]) + self.img = cv2.cvtColor(self.img, cv2.COLOR_RGB2BGR) + self.traktor.set_mode(self.mode(self.recognition(self.img))) + + def mode(self,mode): + self.mode_value = mode + if self.mode_value in [0, 1, 2, 3]: + return 0 + elif self.mode_value in [1, 3, 5, 7]: + return 1 + elif self.mode_value in [0, 1, 4, 5]: + return 2 + elif self.mode_value in [8]: + return 3