import gradio as gr
from PIL import Image
import numpy as np

# Load the pre-trained model from Hugging Face
model = gr.load("models/dima806/indian_food_image_detection")

def classify_image(image):
    # Ensure the image is in the correct format for the model
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    result = model(image)
    
    # Parse the result to the format expected by Gradio's label component
    predictions = result[0]  # Adjust this line based on the actual output structure
    formatted_result = {pred["label"]: pred["score"] for pred in predictions}

    return formatted_result

# Create a Gradio interface with the custom title
iface = gr.Interface(
    fn=classify_image, 
    inputs=gr.Image(type="numpy", label="Upload an image"), 
    outputs=gr.Label(label="Prediction"), 
    title="CAPSTONE",
    description="Upload an image of Indian food to detect what it is."
)

# Launch the Gradio interface
iface.launch()