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

def image_mod(image):
    # img = Image.fromarray(image['composite'])
    

    img = image['composite']
    test_img = np.array(img).reshape(1, test_img.shape[0], test_img.shape[1], 1)
    print(test_img.shape)
    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)