Merge pull request #15 from norkator/Feature/gpu-memory-problem-fix

Feature/gpu memory problem fix

Author: norkator

libraries/fast_srgan/infer_oi.py

@@ -1,5 +1,7 @@
 from argparse import ArgumentParser
+import tensorflow as tf
 from tensorflow import keras
+from keras import backend as kb
 import numpy as np
 import cv2
 import os
@@ -11,15 +13,22 @@
 # Model path
 model_path_file_name = os.getcwd() + '/libraries/fast_srgan/models/generator.h5'
 
-# Load model to memory
-# Change model input shape to accept all size inputs
-model = keras.models.load_model(model_path_file_name, compile=False)
-inputs = keras.Input((None, None, 3))
-output = model(inputs)
-model = keras.models.Model(inputs, output)
+# Set Keras TensorFlow session config
+config = tf.compat.v1.ConfigProto()
+config.gpu_options.per_process_gpu_memory_fraction = 0.9  # Use 90%
+config.gpu_options.allow_growth = True
+tf_session = tf.compat.v1.Session(config=config)
+tf.compat.v1.keras.backend.set_session(tf_session)
 
 
 def process_super_resolution_images(sr_image_objects):
+    # Load model to memory
+    # Change model input shape to accept all size inputs
+    model = keras.models.load_model(model_path_file_name, compile=False)
+    inputs = keras.Input((None, None, 3))
+    output = model(inputs)
+    model = keras.models.Model(inputs, output)
+
     # Loop over all images
     # Input and output image is full path + filename including extension
     for sr_image_object in sr_image_objects:
@@ -36,7 +45,8 @@ def process_super_resolution_images(sr_image_objects):
             print('Sr processing img size: ' + str(original_h) + ':' + str(original_w))
 
             # Check if image is not too big
-            if original_w < 1200 and original_h < 1200:
+            # Original size was 1200 but testing with smaller
+            if original_w < 1000 and original_h < 1000:
                 # Convert to RGB (opencv uses BGR as default)
                 low_res = cv2.cvtColor(low_res, cv2.COLOR_BGR2RGB)
 
@@ -57,11 +67,19 @@ def process_super_resolution_images(sr_image_objects):
 
                 # Save sr image data to object
                 sr_image_object.set_sr_image_data(sr)
+
+                # Clear sr object
+                sr = None
             else:
                 # Save original image
                 sr_image_object.set_sr_image_data(low_res)
 
         except Exception as e:
             print(e)
 
+    # Clear model
+    model = None
+    tf.compat.v1.reset_default_graph()  # Try free memory from Tensorflow
+    kb.clear_session()  # Clear Keras session
+
     return sr_image_objects