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

# Load the model
model = tf.keras.models.load_model('mnist_trained_model_2(acc-97%).h5')

# Classification prediction function
labels = [0,1,2,3,4,5,6,7,8,9]
def classify_image(image):
  prediction = model.predict(image.reshape(-1,784)/255).flatten()
  confidences = {labels[i]: float(prediction[i]) for i in range(10)}
  return confidences
    
label = gr.outputs.Label(num_top_classes=3)
interface = gr.Interface(fn=classify_image, 
             inputs="sketchpad",
             outputs=label,
             capture_session="True",
             title = "Digits Classification",
             description = "This is a machine learning model which can classify the hand written digits from 0-9.")
# Launch
interface.launch(inline = False)