from PIL import Image
import matplotlib.pyplot as plt
import gradio as gr
from transformers import pipeline
from transformers import AutoModelForImageClassification, AutoImageProcessor

image_processor = AutoImageProcessor.from_pretrained("./Mymodel/")
model = AutoModelForImageClassification.from_pretrained("./Mymodel/")

food_info = {
    'Breakfast Burrito': {
        'Description': "Wrap up your mornings with a burst of flavor – the breakfast burrito, where every bite is a sunrise in your mouth!",
        'Calories and Health Info': "The calorie count can vary depending on ingredients, but incorporating whole-grain tortillas, lean proteins, and veggies can make it a balanced and nutritious choice for a hearty breakfast."
    },
    'Caesar Salad': {
        'Description': "Crispy, creamy, and utterly satisfying – the Caesar Salad, where freshness meets indulgence in every crunchy bite!",
        'Calories and Health Info': "While delicious, be mindful of the dressing and croutons. Opt for a lighter dressing and whole-grain croutons to keep it healthier. Loaded with vitamins and greens, it's a generally nutritious option."
    },
    'Chicken Quesadilla': {
        'Description': "Cheesy, savory, and filled with zest – the chicken quesadilla, because life is better when wrapped in a warm tortilla!",
        'Calories and Health Info': "Moderation is key. Choose lean chicken and load up on veggies to make it a protein-packed, flavorful treat. Use whole-grain tortillas for added nutritional value."
    },
    'Club Sandwich': {
        'Description': "Layers of delight stacked high – the club sandwich, where every layer tells a story of taste and texture!",
        'Calories and Health Info': "Opt for whole-grain bread, lean proteins, and plenty of veggies. The club sandwich can be a balanced choice with a mix of carbohydrates, proteins, and essential nutrients."
    },
    'Donuts': {
        'Description': "Ring-shaped bliss, glazed with happiness – donuts, where every bite is a moment of pure joy!",
        'Calories and Health Info': "A sweet indulgence best enjoyed in moderation. High in sugars and fats, it's a treat for special occasions rather than an everyday choice."
    },
    'Fish And Chips': {
        'Description': "Crispy ocean delights paired with golden perfection – fish and chips, a symphony of crunch and flavor!",
        'Calories and Health Info': "While delicious, it can be high in calories and fats. Opt for baked or grilled fish and consider sweet potato fries for a healthier twist."
    },
    'Hamburger': {
        'Description': "Juicy, savory, and the epitome of handheld happiness – the hamburger, where every bite is a taste of tradition!",
        'Calories and Health Info': "Choose lean meats, whole-grain buns, and load up on veggies for a more balanced option. Moderation is key to enjoying this classic."
    },
    'Pizza': {
        'Description': "A slice of heaven in every bite – pizza, where the melding of cheese, sauce, and crust creates an edible masterpiece!",
        'Calories and Health Info': "Balance is key. Opt for thin crust, load up on veggies, and consider lean protein toppings for a more nutritious pizza experience."
    },
    'Samosa': {
        'Description': "A crispy pocket of delight – samosa, where every fold tells a tale of spices and flavors!",
        'Calories and Health Info': "While a tasty snack, samosas are often fried. Enjoy in moderation and consider baked versions with a variety of fillings for a healthier option."
    },
    'Waffles': {
        'Description': "Crisp on the outside, fluffy on the inside – waffles, turning breakfast into a golden affair!",
        'Calories and Health Info': "Opt for whole-grain or alternative flours, and top with fresh fruits or yogurt for added nutrition. Enjoying waffles in moderation can be part of a balanced breakfast."
    }
}

title = "Foodie 🍕"
description = " Image classification model capable of accurately predicting 10 different foods."
article = "Can predict the following 10 classes: Breakfast Burrito, Caesar Salad, Chicken Quesadilla, Club Sandwich, Donuts, Fish And Chips, Hamburger, Pizza, Samosa and Waffles."

import os

folder_path = "Images" 
example_list = []
if os.path.exists(folder_path) and os.path.isdir(folder_path):
    file_paths = [os.path.join(folder_path, file_name) for file_name in os.listdir(folder_path)]
    for file_path in file_paths:
        example_list.append([file_path])

def predict(my_image):
    image = Image.fromarray(my_image.astype('uint8'))

    pipe = pipeline("image-classification", 
                    model=model,
                    feature_extractor=image_processor)
    
    pred = pipe(image)
    
    res = {}
    for i in pred:
        res[i['label'].replace('_', ' ').title()] = round(i['score'])
    return res, food_info[pred[0]['label'].replace('_', ' ').title()]['Description'],food_info[pred[0]['label'].replace('_', ' ').title()]['Calories and Health Info']

iface = gr.Interface(fn=predict,
                    inputs='image',
                    outputs=[gr.Label(num_top_classes=5, label="Predictions"), 
                             gr.Text(label='Description'),
                             gr.Text(label='Calories and Health Info')],
                             title=title,
                             description=description,
                             article=article,
                             examples=example_list)
iface.launch()