load_util.py
966 Bytes
import os
import cv2
import yaml
import tensorflow as tf
def load_config(config_path):
with open(config_path, 'r') as cf:
config_obj = yaml.load(cf, Loader=yaml.FullLoader)
print(config_obj)
return config_obj
def load_argue_model(config):
cls_argue_path = config['MODEL']['CLS_ARGUE']
with tf.Graph().as_default():
if os.path.isfile(cls_argue_path):
print('Model filename: %s' % cls_argue_path)
with tf.gfile.GFile(cls_argue_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
x = tf.get_default_graph().get_tensor_by_name("x_batch:0")
output = tf.get_default_graph().get_tensor_by_name("output/BiasAdd:0")
config = tf.ConfigProto()
config.gpu_options.allow_growth = False
sess = tf.Session(config=config)
return x, output, sess