projekt_sztuczna_int/cnn.py

34 lines
970 B
Python
Raw Permalink Normal View History

2021-06-19 11:28:37 +02:00
#wersja 1.061, na naowszych nie będzie działać
from fastai.vision import *
import os, random
class neuralNetwork:
def __init__(self):
self.trash_types = {
4: 'plastic',
1:'glass',
5:'municipal',
0:'municipal',
2:'municipal',
3:'paper',
}
self.lerner = load_learner('D:\\Pobieranie\\sztuczna_int-main\\sztuczna_int-main')
def getPredict(self):
rand = random.choice(os.listdir("D:\\Pobieranie\\sztuczna_int-main\\sztuczna_int-main\\Test"))
path = "D:\\Pobieranie\\sztuczna_int-main\\sztuczna_int-main\\Test\\" + rand
print(rand)
img = open_image(path)
interpret=self.lerner.predict(img)[2]
max_idxs = np.asarray(np.argmax(interpret[0]))
print(self.trash_types[self.lerner.predict(img)[1].item()])
print(interpret.tolist())
return self.trash_types[self.lerner.predict(img)[1].item()]