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