import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing.image import img_to_array
from keras.datasets import mnist

(X_train, y_train), (X_test, y_test) = mnist.load_data()

X_train = X_train/255.0
X_test = X_test/255.0

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(units=128, activation='relu'),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(units=10,activation="softmax")
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy']
)
history = model.fit(X_train,y_train,epochs=10,validation_data=(X_test,y_test))

def predict(img):
    x = img_to_array(img)
    x = np.expand_dims(x,axis=0)
    target = model.predict(x)
    target = np.argmax(target)
    return target

demo = gr.Interface(fn=predict, 
             inputs="sketchpad",
             outputs="number",
             live=True, streaming=True)

demo.launch()