import tensorflow as tf
import os
import numpy as np
import cv2

from PIL import Image
from PIL import ImageOps

from tensorflow.keras import datasets, layers, models

LEARNING_RATE = 1e-4
DATADIR = "" # path to the dataset directory
CATEGORIES = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
MODEL_PATH = "" # path to the tflite model

if __name__ == '__main__': 
  test_data = []

  for category in CATEGORIES:
    path = os.path.join(DATADIR, category)
    class_num = CATEGORIES.index(category)
    for img in os.listdir(path):
      img_array = cv2.imread(os.path.join(path, img), cv2.IMREAD_GRAYSCALE)
      test_data.append([img_array, class_num, img])

  interpreter = tf.lite.Interpreter(model_path=MODEL_PATH)
  interpreter.allocate_tensors()
  
  input_details = interpreter.get_input_details()
  output_details = interpreter.get_output_details()
  
  labels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  i=0
  ok = 0
  
  for image in test_data:
      i+=1
      input_data = image[0]      
      input_data = np.array(input_data, dtype=np.float32).reshape(-1, 28, 28, 1)
      input_data = input_data/255.0
      
      interpreter.set_tensor(input_details[0]['index'], input_data)
      interpreter.invoke()
      
      output_data = interpreter.get_tensor(output_details[0]['index'])
      results = np.squeeze(output_data)

      top_k = results.argsort()[-5:][::-1]

      if top_k[0] == image[1]:
        ok += 1
      else:
        print('Predicted: %d, Actual: %d, File: %s' % (labels[top_k[0]], image[1], image[2]))
        
  print(ok/i)
