Working trash node generation image based

Logic/TrashRecognition/ImageClassification.py

@@ -72,13 +72,13 @@ def load_labels(label_file):
     label.append(l.rstrip())
   return label
 
-def classify(model="Model/graph.pb", 
-  label_file="Model/retrained_labels.txt",
+def classify(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="input", #"input",
+  input_layer="Mul", #"input",
   output_layer="final_result"): # "InceptionV3/Predictions/Reshape_1"):
   """Returns list of tuples consisting of name of file, category and certainity (0 - 1)"""
   graph = load_graph(model_file)
@@ -112,7 +112,7 @@ def classify(model="Model/graph.pb",
 
 if __name__ == "__main__":
   model_file = "Model/graph.pb"
-  label_file = "Model/retrained_labels.txt"
+  label_file = "Model/graph_labels.txt"
   input_height = 299
   input_width = 299
   input_mean = 128
@@ -165,7 +165,7 @@ if __name__ == "__main__":
   if args.output_layer:
     output_layer = args.output_layer
   
-  classify(model=model_file, label_file=label_file, input_height=input_height, input_width=input_width,
+  classify(model_file=model_file, label_file=label_file, input_height=input_height, input_width=input_width,
     input_mean=input_mean, input_std=input_std, input_layer=input_layer, output_layer=output_layer)
   
     # for i in top_k:

UI/grid.py

@@ -3,8 +3,11 @@ import numpy as np
 import random as rd
 from os import listdir
 from os.path import isfile, join
+from Logic.TrashRecognition.ImageClassification import classify
 
-
+# MODULE LEVEL VARIABLES
+trash_files = classify()
+########################
 
 
 class Grid:
@@ -91,16 +94,24 @@ class House:
         self.empty = True
         self.trash = None
         self.trash_file = None
+
             
     def find_trash_file(self, trash):
-        trash_files_list = []
+        # trash_files_list = []
 
-        file_names = [f for f in listdir("Images\\TestImages") if isfile(join("Images\\TestImages", f))]
-        #filter names
-        for f in file_names:
-            if trash[2] in f:
-                trash_files_list.append(f)
+        # file_names = [f for f in listdir("Images\\TestImages") if isfile(join("Images\\TestImages", f))]
+        # #filter names
+        # for f in file_names:
+        #     if trash[2] in f:
+        #         trash_files_list.append(f)
+
+        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]