forked from s444399/AI
44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
|
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):
|
||
|
|
image = cv2.imread(path)
|
||
|
|
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")
|
||
|
|
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 = []
|
||
|
|
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([classes[ids] for ids in class_ids])
|
||
|
|
x = [1, 0, 0, 0, 0]
|
||
|
|
numpy.random.shuffle(x)
|
||
|
|
if x[0]==1:
|
||
|
|
print("Zwykła")
|
||
|
|
elif x[1]==1:
|
||
|
|
print("Kruchy")
|
||
|
|
elif x[2]==1:
|
||
|
|
print("Łatwopalny")
|
||
|
|
elif x[3]==1:
|
||
|
|
print("Radioaktywny")
|
||
|
|
elif x[4]==1:
|
||
|
|
print("Niebezpieczny")
|
||
|
|
return [list(x)]
|