import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import gradio as gr
import os

print("Loading model...")
model = load_model("food17_inception.keras")
print("Model loaded.")

food_items = [
    "spring_rolls",
    "ice_cream",
    "samosa",
    "strawberry_shortcake",
    "nachos",
    "falafel",
    "tacos",
    "onion_rings",
    "ravioli",
    "hot_dog",
    "apple_pie",
    "pizza",
    "donuts",
    "french_fries",
    "waffles",
    "chocolate_cake",
    "pancakes",
]

# Sort food_items once
food_items.sort()

HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "food17-flagged")

def predict(img):
    print("Predicting...")
    # Resize image to the target size the model expects
    img = img.resize((299, 299))  # Replace with your model's expected input size

    # Convert the image to a numpy array and preprocess it
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array /= 255.0  # Assuming the model expects values between 0 and 1

    # Make prediction
    pred = model.predict(img_array)
    index = np.argmax(pred)
    pred_value = food_items[index]

    probability = pred[0][index]

    print(f"Prediction: {pred_value} ({probability*100:.2f}%)")
    return f"{pred_value} ({probability*100:.2f}%)"

# Define the Gradio interface
interface = gr.Interface(
    fn=predict,
    title="Food Image Classifier",
    description="Upload an image of food, and the model will predict the food label along with the probability. After Prediction Please provide your feedback by selecting one of the flagging options.",
    inputs=gr.inputs.Image(type="pil", label="Select an image of food"),
    outputs="text",
    allow_flagging="manual",
    flagging_options=["Incorrect", "Correct", "Correct but Low Confidence"],
    flagging_callback=hf_writer,
    flagging_dir="flagged_data",
)

print("Launching Gradio interface...")
interface.launch(debug=True)