import numpy
import cv2
import tensorflow as tf
from tensorflow.tools.graph_transforms import TransformGraph

if __name__ == '__main__':
    batch_size = 4
    input_shape = [batch_size, 32, 32, 1]
    features = tf.placeholder(tf.float32, input_shape, name='input')

    features0 = tf.layers.conv2d(inputs=features, filters=8, kernel_size=[3, 3], padding='same')
    # CHANGE THIS TO True to TRIGGER CRASH
    if True:
        features1 = tf.split(features0, num_or_size_splits=2, axis=3)[0]
        features2 = tf.split(features0, num_or_size_splits=2, axis=3)[1]
    else:
        features1, features2 = tf.split(features0, num_or_size_splits=2, axis=3)

    in1 = tf.layers.conv2d(inputs=features1, filters=8, kernel_size=[3, 3], strides=(2, 2), padding='same')
    in2 = tf.layers.conv2d(inputs=features2, filters=8, kernel_size=[3, 3], strides=(2, 2), padding='same')

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        constant_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(),
                                                                      ['conv2d_1/BiasAdd', 'conv2d_2/BiasAdd'])
        tf.train.write_graph(constant_graph, "", "graph_final.pb", as_text=False)

    # export
    with tf.gfile.FastGFile("graph_final.pb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        graph_def = TransformGraph(graph_def, ['input'], ['conv2d_1/BiasAdd', 'conv2d_2/BiasAdd'],
                                   ['strip_unused_nodes'])
        with tf.gfile.FastGFile('saved_model.pb', 'wb') as f:
            f.write(graph_def.SerializeToString())

    # read model
    cvNet = cv2.dnn.readNetFromTensorflow('./saved_model.pb')
    img = numpy.zeros((32, 32, 1), dtype='uint8')
    cvNet.setInput(cv2.dnn.blobFromImage(img, size=(32, 32), swapRB=False, crop=False).repeat(batch_size, axis=0))
    cvOut = cvNet.forward()
    print len(cvOut)