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")
model1 = from_pretrained_keras("ml-debi/EfficientNetV2S-Food101")

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

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

Welcome to my Hugging Face space, where gastronomy meets technology! This space is dedicated to an image classifier that has been trained on the Food101 dataset. 
At the core of this classifier lies the EfficientNetV2S base model, which has been fine-tuned to enhance its ability to identify a broad spectrum of food items with great accuracy.
While it may not be perfect, it’s a nice tool for culinary exploration. Enjoy your journey through the world of food! 🍽️
"""

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

#This model is based on the EfficientNetB0 architecture and was trained on augmented data, providing improved generalization.
#"""
examples = [["./examples/club_sandwich.jpg"], ["./examples/edamame.jpg"], ["./examples/eggs_benedict.jpg"]]

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)[0]
    print("model prediction", prediction)
    confidences = {model1.config['id2label'][str(i)]: float(prediction[i]) for i in range(101)}
    return confidences

iface = gr.Interface(
    fn=predict,
    inputs=[gr.Image()],
    outputs=[gr.Label(num_top_classes=5)],
    title="Food Vision Mini Project",
    description=f"{model1_info}\n",
    examples=examples
)

iface.launch()