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

# Configuration
HEIGHT, WIDTH = 224, 224
NUM_CLASSES = 6
LABELS = ["Burger King", "KFC", "McDonalds", "Other", "Starbucks", "Subway"]
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_model2.h5')

# Loading trained model
# model = tf.keras.models.load_model('best_model2.h5')

def classify_image(inp):
    np.random.seed(143)

    # Preprocess input
    inp = inp.reshape((-1, HEIGHT, WIDTH, 3))
    inp = tf.keras.applications.nasnet.preprocess_input(inp)

    # Prediction
    prediction = model.predict(inp)
    # Build a dict of label:confidence
    return {LABELS[i]: float(f"{prediction[0][i]:.6f}") for i in range(NUM_CLASSES)}

# Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(
        label="Input Image",
        sources="upload",      # or "sketchpad", "webcam"
        type="numpy",         # pass as numpy array to your function
        height=HEIGHT,        # set display height :contentReference[oaicite:0]{index=0}
        width=WIDTH           # set display width :contentReference[oaicite:1]{index=1}
    ),
    outputs=gr.Label(num_top_classes=4),
    title="Brand Logo Detection"
)

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