#!/usr/bin/env python3
import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import numpy
import cv2
import tensorflow as tf
from tensorflow.keras.layers import Cropping2D, Input, Conv2D
from tensorflow.tools.graph_transforms import TransformGraph

if __name__ == '__main__':
  print(cv2.__version__)
  batch_size = 4
  input_shape = [batch_size, 32, 32, 2]
  features = Input(input_shape[1:], name='input')
  features0 = Conv2D(filters=8, kernel_size=[3, 3], padding='same')(features)
  features0 = Cropping2D(((0, 1), (0, 1)))(features0)

  with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      res = sess.run(features0, {features: numpy.zeros(input_shape)})
      print(res.shape)
      constant_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(),
                                                                    ['cropping2d/strided_slice'])
      tf.train.write_graph(constant_graph, "", "graph_final.pb", as_text=False)

  # export
  with tf.gfile.FastGFile("graph_final.pb", 'rb') as f:
      graph_def = tf.GraphDef()
      graph_def.ParseFromString(f.read())
      graph_def = TransformGraph(graph_def, ['input'], ['cropping2d/strided_slice'],
                                 ['strip_unused_nodes', 'fold_constants'])
      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((batch_size, 2, 32, 32), dtype='float32')
  cvNet.setInput(img)
  cvOut = cvNet.forward()
  print(cvOut.shape)