import tensorflow as tf
import cv2
import numpy as np
import argparse


def load_graph(path):
    with tf.gfile.FastGFile(path, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(graph_def, name='')

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', required = True, help = 'Path to image')
    parser.add_argument('--gdef', required = True, help = 'Path to gdef')
    parser.add_argument('--output', required = True, help = 'Path to output of network')
    args = parser.parse_args()

    load_graph(args.gdef)

    img = cv2.imread(args.image)
    if img is None:
        raise IOError("Failed to open file '{}'".format(args.image))
    net_input = np.expand_dims(img, axis = 0)

    with tf.Session() as sess:
        sigmoid = sess.graph.get_tensor_by_name('inference_sigmoid:0')
        fuse = sess.run(sigmoid, {'inference_input:0': net_input})
        cv2.imwrite(args.output, np.squeeze(fuse)*255)

if __name__ == '__main__':
    main()
