# -*- coding:utf-8 -*-
import os
import numpy as np
import time
import io
from PIL import Image

import tensorflow as tf

test_image_dir = 'data/test'
model_path = "model_quantized.tflite"
chars='23456789abcdefghjkmpqrstuvwxy'

# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
print(str(input_details))
output_details = interpreter.get_output_details()
print(str(output_details))

def img2input(img,width,height):
  tmpe_array=[]
  for i in range(height):
    for j in range(width):
      pixel=img.getpixel((j,i))
      tmpe_array.append((0.3*pixel[0]+0.6*pixel[1]+0.1*pixel[2])/255)
  tmpe_array=np.array(tmpe_array).astype('float32')
  input_array=np.expand_dims(tmpe_array, axis=0)
  return input_array

model_interpreter_time = 0
# 计时
start_time = time.time()
# 遍历文件
file_list = os.listdir(test_image_dir) 
for file in file_list:
    print('-----------------------------')
    full_path = os.path.join(test_image_dir, file)       
    captcha_image = Image.open(full_path)
    image_np_expanded=img2input(captcha_image,100,50)
    
    # 填装数据
    model_interpreter_start_time = time.time()
    interpreter.set_tensor(input_details[0]['index'], image_np_expanded)
    
    # 注意注意，我要调用模型了
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])
    model_interpreter_time += time.time() - model_interpreter_start_time
    codes=''
    for i in output_data[0]:
      codes=codes+chars[i]
    print('codes:{}'.format(codes))

used_time = time.time() - start_time
print('used_time:{}'.format(used_time))
print('model_interpreter_time:{}'.format(model_interpreter_time))