import gradio as gr
from PIL import Image
from transformers import pipeline

def convert_to_rgb(image):
    return image.convert('RGB')


classifier = pipeline("image-classification", model="google/vit-base-patch16-224")


def classify_image(image):
    image = convert_to_rgb(image)
    class_scores = classifier(image)
    highest_probability_class = max(class_scores, key=lambda x: x["score"])
    highest_probability_class = highest_probability_class["label"]
    return highest_probability_class, class_scores


iface = gr.Interface(fn=classify_image, 
                     inputs=gr.Image(type="pil"), 
                     outputs=["text", "json"],
                     live=True, 
                     title="Food Image Classification",
                     description="Classify food items in uploaded images using a pre-trained model.")

iface.launch()