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

# model = tf.keras.models.Sequential([
#     tf.keras.layers.Conv2D(filters=16, kernel_size=3, strides=1, padding='same', activation='relu', input_shape=(28, 28, 1)),
#     tf.keras.layers.Conv2D(filters=16, kernel_size=3, strides=1, padding='same', activation='relu'),
#     tf.keras.layers.BatchNormalization(),
#     tf.keras.layers.Conv2D(filters=32, kernel_size=3, strides=1, padding='same', activation='relu'),
#     tf.keras.layers.Conv2D(filters=32, kernel_size=3, strides=1, padding='same', activation='relu'),
#     tf.keras.layers.BatchNormalization(),
#     tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=2, padding='same', activation='relu'),
#     tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=2, padding='same', activation='relu'),
#     tf.keras.layers.BatchNormalization(),
#     tf.keras.layers.GlobalAveragePooling2D(),
#     tf.keras.layers.Dense(10, activation='softmax')
# ])

# model.compile(optimizer=tf.keras.optimizers.Adam(),
#              loss=tf.keras.losses.CategoricalCrossentropy(),
#              metrics=[tf.keras.metrics.MeanSquaredError(), tf.keras.metrics.AUC(), tf.keras.metrics.CategoricalAccuracy()])

# model.load_model("my_model.keras")


def image_mod(image):
    
    # img = Image.fromarray(image['composite'])
    model = tf.keras.models.load_model('weights_1.h5')
    test_img = np.array(image['composite']).reshape(1, 28, 28, 1)
    # test_img = cv2.resize(np.array(image['composite']), (28, 28, 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)