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


model = tf.keras.models.Sequential([
    tf.keras.layers.Input(shape=(28, 28, 1)),
    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 = tf.keras.models.load_model('my_model.h5', compile=False)


def classify_image(image):
    if len(np.array(image).shape) == 3:
      image = tf.image.rgb_to_grayscale(image)
    # image = np.array(image['composite'])[:, :, 3] * 255
    # image = image[..., np.newaxis]
    image_tensor = tf.convert_to_tensor(image)
    image_tensor = tf.image.resize(image_tensor, (28, 28)),
    image_tensor = tf.cast(image_tensor, tf.float32)
    image_tensor = image_tensor / 255.0
    prediction = model.predict(image_tensor)
    print(prediction)
    prediction_label = str(prediction.argmax())
    return prediction_label


title = "Draw to Search"
description = "Using the power of AI to detect the number you draw!"
article = "for source code you can visit [my github](https://github.com/mralamdari)"

example_list = [["examples/" + example] for example in os.listdir("examples")]

interface = gr.Interface(fn=classify_image,
                          inputs=gr.Image(type="pil"),
                          # inputs=gr.Sketchpad(),
                          # outputs=gr.Label(num_top_classes=3, label="Predictions"),
                          outputs='text',
                          examples=example_list, 
                          title=title,
                          description=description,
                          article=article)
                        
interface.launch()