import gradio as gr
from transformers import pipeline

# Load the image classification pipeline from Hugging Face Transformers
pipe = pipeline("image-classification", model="heisenberg3376/vit-base-food-items-v1")

# Define the Gradio interface function
def classify_image(input_image):
    # Perform classification on the input image
    results = pipe(input_image)
    
    # Prepare the output string with all predictions
    confidences = {result['label']: float(result['score']) for result in results}
    
    # Return the concatenated string of predictions
    return confidences

# Create a Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Label(num_top_classes=5),
    title="Image Classification",
    description="Classify food items in images using heisenberg3376/vit-base-food-items-v1"
)

# Launch the Gradio interface
iface.launch()