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)























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')

def classify_image(image):
    if len(np.array(image).shape) == 3:
        image = tf.image.rgb_to_grayscale(image)
    image_tensor = tf.convert_to_tensor(image)
    image_tensor = tf.cast(image_tensor, tf.float32)
    image_tensor = tf.expand_dims(image_tensor, 0)
    image_tensor = image_tensor / 255.0
    prediction = model.predict(image_tensor)
    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"),
                          outputs=gr.Label(num_top_classes=3, label="Predictions"),
                          examples=example_list, 
                          title=title,
                          description=description,
                          article=article)
                        
interface.launch()