#!/usr/bin/env python

###############################################################################
#
# Copyright 2017 - 2019 NXP
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
#
###############################################################################

# DESCRIPTION:        Generate the TFLite model
# INPUT/S:            Saved model
# OUTPUT/S:           TFLite model

import tensorflow as tf
import os
import glob

###################################################################
# Main function
###################################################################    

if __name__ == "__main__":
    
    print ("="*73)
    print ("="*32," Start ","="*32)
    print ("="*73)
    
    # Path to saved model directory
    try:
        os.chdir("../Saved_Model")
        export_dir = os.getcwd()
    except:
        print ("\n*** saved model directory is not exist ***")
        
    # Generating TFLite model 
    converter = tf.lite.TFLiteConverter.from_saved_model(export_dir)                                       
    tflite_model = converter.convert()
    
    # Path to save TFLite model
    try:
        os.chdir("../TFlite_Model")
        for file_ in glob.glob("*.xls"):
            os.remove(file_)
    except:
        os.mkdir("../TFlite_Model")
        os.chdir("../TFlite_Model")
        
    open("converted_model_new.tflite", "wb").write(tflite_model)

    print ("="*73)
    print ("="*33," End ","="*33)
    print ("="*73)
