import os
import numpy as np
import gradio as gr
from PIL import Image
import tensorflow as tf

def image_mod(image):
    # img = Image.fromarray(image['composite'])
    model = tf.keras.models.load_model('my_model.keras')
    test_img = np.array(image['composite']).reshape(1, test_img.shape[0], test_img.shape[1], 1)
    prediction = model.predict(test_img)
    pred   = np.argmax(prediction, axis=1)[0]
    return pred


title = "Draw to Search"
description = "Using the power of AI to detect the number you draw!"

demo = gr.Interface(
  fn=image_mod, 
  inputs='sketchpad',
  outputs='text',
  title=title,
  description=description,
  live=True)

demo.launch(share=False)
# demo.launch(debug=True)