pull/35/head
Clemens Neudecker 3 years ago committed by GitHub
parent b7aeb1ae8b
commit 83c60d7b26
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -15,8 +15,8 @@ environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
stderr = sys.stderr stderr = sys.stderr
sys.stderr = open(devnull, 'w') sys.stderr = open(devnull, 'w')
import tensorflow as tf import tensorflow as tf
from tensorflow.keras.models import load_model, Model from tensorflow.keras.models import load_model
from tensorflow.python.keras import backend as K from tensorflow.python.keras import backend as tensorflow_backend
sys.stderr = stderr sys.stderr = stderr
@ -44,10 +44,10 @@ class SbbBinarizer:
config.gpu_options.allow_growth = True config.gpu_options.allow_growth = True
self.session = tf.compat.v1.Session(config=config) # tf.InteractiveSession() self.session = tf.compat.v1.Session(config=config) # tf.InteractiveSession()
K.set_session(self.session) tensorflow_backend.set_session(self.session)
def end_session(self): def end_session(self):
K.clear_session() tensorflow_backend.clear_session()
self.session.close() self.session.close()
del self.session del self.session
@ -59,7 +59,7 @@ class SbbBinarizer:
return model, model_height, model_width, n_classes return model, model_height, model_width, n_classes
def predict(self, model_in, img, use_patches): def predict(self, model_in, img, use_patches):
K.set_session(self.session) tensorflow_backend.set_session(self.session)
model, model_height, model_width, n_classes = model_in model, model_height, model_width, n_classes = model_in
img_org_h = img.shape[0] img_org_h = img.shape[0]

Loading…
Cancel
Save