From b4765ba08f0fa54e637ecd9c642df6a156414129 Mon Sep 17 00:00:00 2001
From: AlexeyAB <kikots@mail.ru>
Date: Mon, 20 Jul 2020 16:28:23 +0300
Subject: [PATCH] Resize the predictions to the original image size

---
 tf/run_onnx.py |   5 +--
 tf/run_pb.py   | 105 +++++++++++++++++++++++--------------------------
 2 files changed, 52 insertions(+), 58 deletions(-)

diff --git a/tf/run_onnx.py b/tf/run_onnx.py
index 6d8e5e6..b20421a 100644
--- a/tf/run_onnx.py
+++ b/tf/run_onnx.py
@@ -52,9 +52,7 @@ def run(input_path, output_path, model_path):
                 resize_method="upper_bound",
                 image_interpolation_method=cv2.INTER_CUBIC,
             )
-
-    #normalize_image = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-
+    
     def compose2(f1, f2):
         return lambda x: f2(f1(x))
 
@@ -81,6 +79,7 @@ def compose2(f1, f2):
         # compute
         output = tf_rep.run(img_input.reshape(1, 3, 384, 384))
         prediction = np.array(output).reshape(384, 384)
+        prediction = cv2.resize(prediction, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)
        
         # output
         filename = os.path.join(
diff --git a/tf/run_pb.py b/tf/run_pb.py
index a86d926..c0a7ba3 100644
--- a/tf/run_pb.py
+++ b/tf/run_pb.py
@@ -24,81 +24,76 @@ def run(input_path, output_path, model_path):
     if gpus:
       try:
         for gpu in gpus:
-          tf.config.experimental.set_memory_growth(gpu, True)
+          #tf.config.experimental.set_memory_growth(gpu, True)
           tf.config.experimental.set_virtual_device_configuration(gpu,
             [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4000)])
       except RuntimeError as e:
         print(e)
 
-    # select device
-    with tf.device('/gpu:0'):
-
-        # load network
-        graph_def = tf.compat.v1.GraphDef()
-        with tf.io.gfile.GFile(model_path, 'rb') as f:
-            graph_def.ParseFromString(f.read())
-            tf.import_graph_def(graph_def, name='')
-
-        #output_layer = '1195:0'
-        #input_node = '0:0'
-        model_operations = tf.compat.v1.get_default_graph().get_operations()
-        input_node = '0:0'
-        output_layer = model_operations[len(model_operations) - 1].name + ':0'
-        print("Last layer name: ", output_layer)
-
-        resize_image = Resize(
-                    384,
-                    384,
-                    resize_target=None,
-                    keep_aspect_ratio=False,
-                    ensure_multiple_of=32,
-                    resize_method="upper_bound",
-                    image_interpolation_method=cv2.INTER_CUBIC,
-                )
-
-        #normalize_image = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-
-        def compose2(f1, f2):
-            return lambda x: f2(f1(x))
+    # load network
+    graph_def = tf.compat.v1.GraphDef()
+    with tf.io.gfile.GFile(model_path, 'rb') as f:
+        graph_def.ParseFromString(f.read())
+        tf.import_graph_def(graph_def, name='')
+
+    
+    model_operations = tf.compat.v1.get_default_graph().get_operations()
+    input_node = '0:0'
+    output_layer = model_operations[len(model_operations) - 1].name + ':0'
+    print("Last layer name: ", output_layer)
+
+    resize_image = Resize(
+                384,
+                384,
+                resize_target=None,
+                keep_aspect_ratio=False,
+                ensure_multiple_of=32,
+                resize_method="upper_bound",
+                image_interpolation_method=cv2.INTER_CUBIC,
+            )
+    
+    def compose2(f1, f2):
+        return lambda x: f2(f1(x))
 
-        transform = compose2(resize_image, PrepareForNet())
+    transform = compose2(resize_image, PrepareForNet())
 
-        # get input
-        img_names = glob.glob(os.path.join(input_path, "*"))
-        num_images = len(img_names)
+    # get input
+    img_names = glob.glob(os.path.join(input_path, "*"))
+    num_images = len(img_names)
 
-        # create output folder
-        os.makedirs(output_path, exist_ok=True)
+    # create output folder
+    os.makedirs(output_path, exist_ok=True)
 
-        print("start processing")
+    print("start processing")
 
-        with tf.compat.v1.Session() as sess:
-          try:
-            # load images
-            for ind, img_name in enumerate(img_names):
+    with tf.compat.v1.Session() as sess:
+      try:
+        # load images
+        for ind, img_name in enumerate(img_names):
 
-                print("  processing {} ({}/{})".format(img_name, ind + 1, num_images))
+            print("  processing {} ({}/{})".format(img_name, ind + 1, num_images))
 
-                # input
-                img = utils.read_image(img_name)
-                img_input = transform({"image": img})["image"]
+            # input
+            img = utils.read_image(img_name)
+            img_input = transform({"image": img})["image"]
 
-                # compute
-                prob_tensor = sess.graph.get_tensor_by_name(output_layer)
-                prediction, = sess.run(prob_tensor, {input_node: [img_input] })
-                prediction = prediction.reshape(384, 384)
-       
+            # compute
+            prob_tensor = sess.graph.get_tensor_by_name(output_layer)
+            prediction, = sess.run(prob_tensor, {input_node: [img_input] })
+            prediction = prediction.reshape(384, 384)
+            prediction = cv2.resize(prediction, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)
+            
             # output
             filename = os.path.join(
                 output_path, os.path.splitext(os.path.basename(img_name))[0]
             )
             utils.write_depth(filename, prediction, bits=2)
 
-          except KeyError:
-            print ("Couldn't find input node: ' + input_node + ' or output layer: " + output_layer + ".")
-            exit(-1)
+      except KeyError:
+        print ("Couldn't find input node: ' + input_node + ' or output layer: " + output_layer + ".")
+        exit(-1)
 
-        print("finished")
+    print("finished")
 
 
 if __name__ == "__main__":