import onnx
from keras.models import *
import keras
import keras2onnx
from train import get_model
import densenet

basemodel, model = get_model(32, 88) # input height, classes number
basemodel.load_weights("models/weights_densenet-32-0.40.h5")
onnx_model = keras2onnx.convert_keras(basemodel, basemodel.name, target_opset=12)
onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = int(1)
onnx_model.graph.input[0].type.tensor_type.shape.dim[1].dim_value = int(1)
onnx_model.graph.input[0].type.tensor_type.shape.dim[2].dim_value = int(32)
onnx_model.graph.input[0].type.tensor_type.shape.dim[3].dim_value = int(280)
onnx_model.graph.output[0].type.tensor_type.shape.dim[0].dim_value = int(1)
onnx_model.graph.node.remove(onnx_model.graph.node[0])
onnx_model.graph.node[0].input[0] = "the_input"
onnx.save_model(onnx_model, "./densenet_ctc.onnx")