AI/neurons.py

58 lines
1.7 KiB
Python
Raw Normal View History

2020-04-28 12:18:53 +02:00
import numpy
import cv2
import argparse
import numpy as np
class Neurons:
def __init__(self):
pass
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 whatIsIt(self, path):
2020-05-08 22:36:01 +02:00
image = cv2.imread(path)
2020-04-28 12:18:53 +02:00
scale = 0.00392
classes = None
with open("yolov3.txt", 'r') as f:
classes = [line.strip() for line in f.readlines()]
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
2020-05-08 22:36:01 +02:00
blob = cv2.dnn.blobFromImage(image, scale, (608, 608), (0, 0, 0), True, crop=False)
2020-04-28 12:18:53 +02:00
net.setInput(blob)
outs = net.forward(self.get_output_layers(net))
class_ids = []
for out in outs:
for detection in out:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
2020-05-08 22:36:01 +02:00
if confidence > 0.2:
2020-04-28 12:18:53 +02:00
class_ids.append(class_id)
2020-05-08 22:36:01 +02:00
y = [classes[ids] for ids in class_ids]
print(y)
x = [0, 0, 0, 0, 0]
if "kruche" in y:
x[1]=1
elif "niebezpieczne" in y:
x[4]=1
elif "biohazard" in y:
x[4]=1
elif "radioaktywne" in y:
x[3]=1
elif "latwopalne" in y:
x[2]=1
2020-05-09 07:28:55 +02:00
else:
x[0]=1
2020-05-08 22:36:01 +02:00
print(x)
2020-04-28 12:18:53 +02:00
if x[0]==1:
print("Zwykła")
2020-05-08 22:36:01 +02:00
if x[1]==1:
2020-04-28 12:18:53 +02:00
print("Kruchy")
2020-05-08 22:36:01 +02:00
if x[2]==1:
2020-04-28 12:18:53 +02:00
print("Łatwopalny")
2020-05-08 22:36:01 +02:00
if x[3]==1:
2020-04-28 12:18:53 +02:00
print("Radioaktywny")
2020-05-08 22:36:01 +02:00
if x[4]==1:
2020-04-28 12:18:53 +02:00
print("Niebezpieczny")
2020-05-08 22:36:01 +02:00
return [list(x)]