#
# (C) Copyright 1996- ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.
#

import os
import numpy as np
import argparse
import keras
import keras2onnx


if __name__ == "__main__":
    """
    Lightweight script to convert a keras model into a TFlite model
    """
    
    parser = argparse.ArgumentParser("Data Augmentation")
    parser.add_argument('keras_model_path', help="Path of the input keras model")
    parser.add_argument('onnx_model_path', help="Path of the output onnx model")
    parser.add_argument("--verify_with", help="Check the model by passing an input numpy path")
    
    args = parser.parse_args()

    # load the keras model
    model = keras.models.load_model(args.keras_model_path)
    model.summary()

    # do the conversion
    onnx_model = keras2onnx.convert_keras(model, model.name)
    
    # write to file
    file = open(args.onnx_model_path, "wb")
    file.write(onnx_model.SerializeToString())
    file.close()

[((23, 13, 23, 57), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({(23, 37, 23, 56): '"""Data Augmentation"""'}, {}), "('Data Augmentation')", False, 'import argparse\n'), ((31, 12, 31, 58), 'keras.models.load_model', 'keras.models.load_model', ({(31, 36, 31, 57): 'args.keras_model_path'}, {}), '(args.keras_model_path)', False, 'import keras\n'), ((35, 17, 35, 60), 'keras2onnx.convert_keras', 'keras2onnx.convert_keras', ({(35, 42, 35, 47): 'model', (35, 49, 35, 59): 'model.name'}, {}), '(model, model.name)', False, 'import keras2onnx\n')]