main.py 文件源码

def load_inference(sess, ckptdir, threshold):
    images = tf.placeholder(tf.float32, shape=[None, IMG_SIZE, IMG_SIZE, 3])
    net = create_model(images, .1)
    sess.run(tf.global_variables_initializer())
    saver = tf.train.Saver(max_to_keep=10)
    if ckptdir and os.path.exists(ckptdir) and not FLAGS.debug:
        checkpoint = tf.train.latest_checkpoint(ckptdir)
        if checkpoint:
            print('Restoring', checkpoint)
            saver.restore(sess, checkpoint)
    return inference(net, threshold), images

###########
# Helpers #
###########