import gradio as gr
import tensorflow as tf
from huggingface_hub import from_pretrained_keras
from tensorflow.keras import mixed_precision

# Load your trained models
model1 = from_pretrained_keras("ml-debi/EfficientNetB0-Food101")
#model2 = from_pretrained_keras("ml-debi/EfficientNetB0-Food101")

with open('classes.txt', 'r') as f:
    classes = [line.strip() for line in f]

# Add information about the models
model1_info = """
### Model 1 Information

This model is based on the EfficientNetB0 architecture and was trained on the Food101 dataset.
"""

#model2_info = """
#### Model 2 Information

#This model is based on the EfficientNetB0 architecture and was trained on augmented data, providing improved generalization.
#"""

def preprocess(image):
    print("before resize", image.shape)
    image = tf.image.resize(image, [224, 224])
    
    image = tf.expand_dims(image, axis=0)
    print("After expanddims", image.shape)
    return image

def predict(image):
    # Choose the model based on the dropdown selection
    #print("---model_selection---", model_selection) #
    #model = model1 if model_selection == "EfficentNetB0 Fine Tune" else model2

    #print(model.summary())
    if mixed_precision.global_policy() == "mixed_float16":
        mixed_precision.set_global_policy(policy="float32")

    image = preprocess(image)
    print(mixed_precision.global_policy())
    prediction = model1.predict(image)
    print("model prediction", prediction)
    predicted_class = classes[int(tf.argmax(prediction, axis=1))]
    confidence = tf.reduce_max(prediction).numpy()
    return predicted_class, confidence

iface = gr.Interface(
    fn=predict,
    inputs=[gr.Image()],
    outputs=[gr.Textbox(label="Predicted Class"), gr.Textbox(label="Confidence")],
    title="Transfer Learning Mini Project",
    description=f"{model1_info}\n",
)

iface.launch()