classification constructor added

classification.py

@@ -3,16 +3,18 @@
 import cv2 as cv
 
 class Classification:
-    def sigmoid(x):
-        return 1 / (1 + np.exp(-x))
+    graph = None
 
-    def load_pb(path_to_pb):
+    def __init__(self, path_to_pb): 
         with tf.io.gfile.GFile(path_to_pb, "rb") as f:
             graph_def = tf.compat.v1.GraphDef()
             graph_def.ParseFromString(f.read())
-            return graph_def
+            self.graph = graph_def
+
+    def sigmoid(x):
+        return 1 / (1 + np.exp(-x))
 
-    def classify(graph, imgPath):
+    def classify(self, img):
         image_mean = 127.5
         image_std = 127.5
         image_size_x = 299
@@ -22,11 +24,10 @@ def classify(graph, imgPath):
         input_node = 'input:0'
 
         with tf.compat.v1.Session() as sess:
-            tf.import_graph_def(graph, name = "")
+            tf.import_graph_def(self.graph, name = "")
+            g = self.graph
             prob_tensor = sess.graph.get_tensor_by_name(output_layer)
 
-            img = cv.imread(imgPath)
-
             img = cv.resize(img, (image_size_x, image_size_y))
             img = np.expand_dims(img, axis=0)
             img = (img - image_mean) / image_std

evaluate_tf.py

@@ -2,6 +2,7 @@
 import numpy as np
 import argparse
 from classification import Classification
+import cv2 as cv
 
 def main():
     parser = argparse.ArgumentParser()
@@ -11,7 +12,7 @@ def main():
     parser.add_argument('--rightVals', help='.txt val path', default='val.txt')
     argv = parser.parse_args()
 
-    graph = Classification.load_pb(argv.graph)
+    graph = Classification(argv.graph)
 
     with open(argv.labels, 'rt') as f:
         labels = f.read().strip().split('\n')  
@@ -26,8 +27,9 @@ def main():
 
     for i in range(size):
         imgPath = argv.databasePath + vals[i].rsplit(' ')[0]
+        img = cv.imread(imgPath)
         val = int(vals[i].rsplit(' ')[1])
-        results = Classification.classify(graph, imgPath)
+        results = graph.classify(img)
 
         if val == (results[0] - 1):
             top_1 += 1