Merge branch 'master' of https://git.wmi.amu.edu.pl/s434720/2019SZI-Projekt

Logic/TrashRecognition/ImageClassification.py

@@ -72,10 +72,47 @@ def load_labels(label_file):
     label.append(l.rstrip())
   return label
 
+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="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)
+
+  files = []
+  for filename in os.listdir('Images/TestImages'):
+    t = read_tensor_from_image_file(
+      f'Images/TestImages/{filename}',
+      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)
+    
+    files.append((filename, labels[top_k[0]], results[top_k[0]]))
+    print(f'{filename}: {labels[top_k[0]]} with {results[top_k[0]] * 100}% certainity')
+  return files
+
 
 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
@@ -128,29 +165,9 @@ if __name__ == "__main__":
   if args.output_layer:
     output_layer = args.output_layer
   
-  graph = load_graph(model_file)
-
-  for filename in os.listdir('Images/TestImages'):
-    t = read_tensor_from_image_file(
-      f'Images/TestImages/{filename}',
-      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'{filename}: {labels[top_k[0]]} with {results[top_k[0]] * 100}% certainity')
+  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:
     #   print(labels[i], results[i])
  

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:
@@ -90,16 +93,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]