import tensorflow as tf

# Convert the model.
converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(
    graph_def_file='tflite/tflite_graph.pb', # Dir of the ".pb" file
    input_arrays=['normalized_input_image_tensor'],
    input_shapes={'normalized_input_image_tensor' : [1, 300, 300,3]}, # Shape of the images of the model you will need to change the 300 and 300
    output_arrays=['TFLite_Detection_PostProcess', 'TFLite_Detection_PostProcess:1', 'TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3'],
)
converter.allow_custom_ops=True

# converter.quantized_input_stats = {'input' : (0., 1.)}  # mean, std_dev (input range is [-1, 1])
# converter.inference_type = tf.int8 # this is the recommended type.
# converter.inference_input_type=tf.uint8 #optional
# converter.inference_output_type=tf.uint8 #optional
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# Save the model.
with open('quantized_model.tflite', 'wb') as f:
  f.write(tflite_model)