import os
import numpy as np
from PIL import Image
import tensorflow as tf
input_mean = 127.5
input_std = 127.5
converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(
	'mobilenet_v1_0.25_128_frozen.pb', # TensorFlow model file
	['input'], # input node names
	['MobilenetV1/Predictions/Reshape_1'], # output node names
	input_shapes={'input': (1, 128, 128, 3)}) # input tensor dimensions
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.float32
def representative_dataset_gen():
	images = []
	for image in os.listdir('dataset'):
		image = os.path.join('dataset', image)
		if os.path.isfile(image):
			image = np.array(Image.open(image).resize((128, 128)))
			image = (image - input_mean) / input_std
			image = image.astype(np.float32)
			image = image.reshape((1, 128, 128, 3))
			images.append(image)
	for image in images:
		yield [image]
converter.representative_dataset = representative_dataset_gen
tflite_quant_model = converter.convert()
with open('mobilenet_v1_0.25_128_quant_int8.tflite', 'wb') as f:
f.write(tflite_quant_model)