import tensorflow as tf
from tensorflow.keras.models import load_model
import gradio as gr
import numpy as np
import cv2
import numpy as np
import os
import sys
from tensorflow import keras
from tensorflow.keras import layers

EPOCHS = 10
IMG_WIDTH = 30
IMG_HEIGHT = 30
NUM_CATEGORIES = 43
TEST_SIZE = 0.4

def get_model():
    """
    Returns a compiled convolutional neural network model. Assume that the
    `input_shape` of the first layer is `(IMG_WIDTH, IMG_HEIGHT, 3)`.
    The output layer should have `NUM_CATEGORIES` units, one for each category.
    """
    # Experimenting with different architectures of a Sequential model and this is so far best for me
    # Accepts input with shape of (30, 30, 3)
    # Convolutional layer with 65 filters and relu activation function
    # Maxpooled by (2, 2) kernel
    # Second convolutional layer with higher number of filter i.e 256 and relu activation function
    # Flattering nd shape and Dense layer with 450 nuerons and relu activation funtion
    # Output layer with 43 nuerons and softmax activation function

    model = tf.keras.Sequential([
        layers.Conv2D(64, 5, input_shape = (30, 30, 3), name = "conv1", activation="relu"),
        layers.MaxPool2D((2, 2), name = "pool1"),
        layers.Conv2D(256, 3, name = "conv2", activation="relu"),
        layers.MaxPool2D((2, 2), name = "pool2"),
        layers.Flatten(),
        layers.Dense(450, activation = "relu", name = "dense1"),
        layers.Dense(NUM_CATEGORIES-1, activation = "softmax", name = "output")
    ])

    # Printing model summary and compiling with adam algorithm, categorical_crossentropy as 43 output neurons
    model.summary()

    model.compile(
        optimizer = "adam",
        loss = "categorical_crossentropy",
        metrics=["accuracy"]
    )

    return model

# model = load_model("best_traffic_model.h5")
# model.compile(optimizer="adam", loss="categorical_crossentropy")

model = get_model()

labels = ["zero","one","two","three","four","five","six","seven","eight","nine","ten","eleven","twelve","thrteen","fourteen","fifteen","sixteen","seventeen","eightteen","nineteen"]

def predict(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    resized = cv2.resize(gray, (64, 64))
    reshaped = resized.reshape(1, 64, 64, 1)
    out = model.predict(reshaped)
    cls = labels[out.argmax()]
    return cls

ui = gr.Interface(
fn=predict,
inputs=gr.Image(type="numpy"),
outputs=gr.Textbox()
)