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_76.h5')


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

    # Ensure input is resized to expected shape
    inp = tf.image.resize(inp, [HEIGHT, WIDTH])
    inp = tf.cast(inp, tf.float32)  # ensure correct dtype for preprocessing
    inp = tf.keras.applications.nasnet.preprocess_input(inp)
    inp = tf.expand_dims(inp, axis=0)  # make batch dimension

    # Prediction
    prediction = model.predict(inp)

    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)