diff --git a/Logic/TrashRecognition/ImageClassification.py b/Logic/TrashRecognition/ImageClassification.py index 88f5591..72b29d3 100644 --- a/Logic/TrashRecognition/ImageClassification.py +++ b/Logic/TrashRecognition/ImageClassification.py @@ -72,7 +72,41 @@ def load_labels(label_file): label.append(l.rstrip()) return label -def classify(model_file="Model/graph.pb", +def classify_file(file_dir="", + model_file="Model/graph.pb", + label_file="Model/graph_labels.txt", + input_height=299, + input_width=299, + input_mean=128, + input_std=128, + input_layer="Mul", #"input", + output_layer="final_result"): + """Returns tuple consisting of name of file, category and certainity (0 - 1)""" + graph = load_graph(model_file) + t = read_tensor_from_image_file( + file_dir, + input_height=input_height, + input_width=input_width, + input_mean=input_mean, + input_std=input_std) + input_name = "import/" + input_layer + output_name = "import/" + output_layer + input_operation = graph.get_operation_by_name(input_name) + output_operation = graph.get_operation_by_name(output_name) + + with tf.Session(graph=graph) as sess: + results = sess.run(output_operation.outputs[0], { + input_operation.outputs[0]: t + }) + results = np.squeeze(results) + + top_k = results.argsort()[-5:][::-1] + labels = load_labels(label_file) + + print(f'{file_dir}: {labels[top_k[0]]} with {results[top_k[0]] * 100}% certainity') + return (file_dir, labels[top_k[0]], results[top_k[0]]) + +def classify_files(model_file="Model/graph.pb", label_file="Model/graph_labels.txt", input_height=299, input_width=299, diff --git a/Logic/TrashRecognition/__pycache__/ImageClassification.cpython-36.pyc b/Logic/TrashRecognition/__pycache__/ImageClassification.cpython-36.pyc index 2e8281f..0abd905 100644 Binary files a/Logic/TrashRecognition/__pycache__/ImageClassification.cpython-36.pyc and b/Logic/TrashRecognition/__pycache__/ImageClassification.cpython-36.pyc differ diff --git a/UI/grid.py b/UI/grid.py index 653a67a..abab26c 100644 --- a/UI/grid.py +++ b/UI/grid.py @@ -3,10 +3,13 @@ import numpy as np import random as rd from os import listdir from os.path import isfile, join -from Logic.TrashRecognition.ImageClassification import classify +# from Logic.TrashRecognition.ImageClassification import classify # MODULE LEVEL VARIABLES -trash_files = classify() +recognized_trash = { + +} +# trash_files = classify() ######################## @@ -73,19 +76,20 @@ class House: GREY = (192,192,192) #define trash - paper = (5,WHITE,"paper") - glass = (6,SKYBLUE,"glass") - metal = (7,GREY,"metal") - plastic = (8,ORANGE,"plastic") - + trash_dict = { + "paper": (5, WHITE, "paper"), + "glass": (6, SKYBLUE, "glass"), + "metal": (7, GREY, "metal"), + "plastic": (8, ORANGE, "plastic") + } #define days of the week - MONDAY = (1, "Monday", paper, metal) - TUESDAY = (2, "Tuesday", glass) - WEDNESDAY=(3, "Wednesday", plastic, metal) - THURSDAY = (4, "Thursday", glass) - FRIDAY = (5, "Friday", paper, metal) - SATURDAY = (6, "Saturday", plastic) - SUNDAY = (7, "Sunday", metal) + MONDAY = (1, "Monday", trash_dict["paper"], trash_dict["metal"]) + TUESDAY = (2, "Tuesday", trash_dict["glass"]) + WEDNESDAY=(3, "Wednesday", trash_dict["plastic"], trash_dict["metal"]) + THURSDAY = (4, "Thursday", trash_dict["glass"]) + FRIDAY = (5, "Friday", trash_dict["paper"], trash_dict["metal"]) + SATURDAY = (6, "Saturday", trash_dict["plastic"]) + SUNDAY = (7, "Sunday", trash_dict["metal"]) DAYS = [MONDAY, TUESDAY, WEDNESDAY, THURSDAY, FRIDAY, SATURDAY, SUNDAY] @@ -95,40 +99,56 @@ class House: self.trash_file = None - def find_trash_file(self, trash): + def find_trash_file(self): # trash_files_list = [] - - # file_names = [f for f in listdir("Images\\TestImages") if isfile(join("Images\\TestImages", f))] + from os.path import sep # culture and os invariant separator + + file_names = [(join(f"Images{sep}TestImages", f)) + for f in listdir(f"Images{sep}TestImages") + if isfile(join(f"Images{sep}TestImages", f))] # #filter names # for f in file_names: # if trash[2] in f: # trash_files_list.append(f) - - trash_files_list = [] + file_name = file_names[rd.randint(0,len(file_names)) - 1] + from Logic.TrashRecognition.ImageClassification import classify_file + if file_name in recognized_trash: + rt = recognized_trash[file_name] + return (file_name, rt[0], rt[1]) + else: + classification = classify_file(file_dir=file_name) + recognized_trash[file_name] = (classification[1], classification[2]) + return classification - # filter names - for f in trash_files: - if trash[2] in f[1]: - trash_files_list.append(f[0]) + # trash_files_list = [] + + # # filter names + # for f in trash_files: + # if trash[2] in f[1]: + # trash_files_list.append(f[0]) - f = rd.randint(0,len(trash_files_list)) - return trash_files_list[f-1] + # f = rd.randint(0,len(trash_files_list)) + # return trash_files_list[f-1] def generate_trash(self): self.empty = False - num = rd.randint(1, 4) - if num == 1: - self.trash = self.paper - self.trash_file = self.find_trash_file(self.trash) - elif num == 2: - self.trash = self.glass - self.trash_file = self.find_trash_file(self.trash) - elif num == 3: - self.trash = self.metal - self.trash_file = self.find_trash_file(self.trash) - elif num == 4: - self.trash = self.plastic - self.trash_file = self.find_trash_file(self.trash) + + classification = self.find_trash_file() + self.trash = self.trash_dict[classification[1]] + self.trash_file = classification[0] + # num = rd.randint(1, 4) + # if num == 1: + # self.trash = self.paper + # self.trash_file = self.find_trash_file(self.trash) + # elif num == 2: + # self.trash = self.glass + # self.trash_file = self.find_trash_file(self.trash) + # elif num == 3: + # self.trash = self.metal + # self.trash_file = self.find_trash_file(self.trash) + # elif num == 4: + # self.trash = self.plastic + # self.trash_file = self.find_trash_file(self.trash) def get_day_of_week(self, d: int): for day in self.DAYS: