import tensorflow as tf
import gradio as gr
import cv2
import numpy as np

def convert_img_to_tensor(image):
   image = cv2.resize(image, (28,28))
   # adding dummy axes to convert to tensor
   image = image[...,np.newaxis]
   image_tensor = tf.convert_to_tensor(image, dtype=tf.float32)
   image_tensor = image_tensor[np.newaxis, ...]
   return image_tensor

def initialize_prediction(image):
   image = convert_img_to_tensor(image)
   model = tf.keras.models.load_model('ACC_98.77.h5')
   labels = {0:'0',1:'1',2:'2',3:'3', 4:'4', 5:'5', 6:'6', 7:'7',8:'8',9:'9'}
   
   predictions = model.predict(image)
   prediction_label = labels[np.argmax(predictions)]
   return prediction_label

gr.Interface(fn=initialize_prediction, 
            inputs="sketchpad", 
            outputs=gr.outputs.Label(label='What the model thinks.')).launch()