keras_to_tflite.py

""" Convert Keras .h5 model to .tflite """
from tensorflow.contrib import lite

def main(model_path: str) -> None:
  
  try:
    converter = lite.TFLiteConverter.from_keras_model_file(model_path)
    model = converter.convert()
  except Exception as e:
    print(e)
    return None
  
  output_path = f"{model_path.split('.')[0]}.tflite" 
  with open(output_path, 'wb') as file:
    file.write( model )


if __name__ == "__main__":
  from argparse import ArgumentParser
  
  parser = ArgumentParser()
  parser.add_argument("-m", "--model", type=str, required=True, help="Path of .h5 model")
  args = parser.parse_args()
  main(args.model)