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

# Configuration
HEIGHT, WIDTH = 224, 224
NUM_CLASSES = 6
LABELS = ["McDonalds", "Burger King", "Subway", "Starbucks", "KFC", "Other"]

from tensorflow_addons.metrics import F1Score
from keras.utils import custom_object_scope

with custom_object_scope({'Addons>F1Score': F1Score}):
    model = tf.keras.models.load_model('best_model.h5')


def classify_image(inp):
    # Resize & preprocess
    inp = tf.image.resize(inp, [HEIGHT, WIDTH])
    inp = tf.cast(inp, tf.float32)
    inp = tf.keras.applications.nasnet.preprocess_input(inp)
    inp = tf.expand_dims(inp, axis=0)

    # Predict
    prediction = model.predict(inp)[0]
    return {LABELS[i]: float(f"{prediction[i]:.6f}") for i in range(NUM_CLASSES)}

example_list = [
    ["Examples/Untitled.png"],
    ["Examples/Untitled2.png"],
    ["Examples/Untitled3.png"],
    ["Examples/Untitled5.png"]
]


iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(
        label="Input Image",
        sources="upload",
        type="numpy",
        height=HEIGHT,
        width=WIDTH
    ),
    outputs=gr.Label(num_top_classes=4),
    title="Brand Logo Detection",
    examples=example_list
)

if __name__ == "__main__":
    iface.launch(debug=False, share=True)