import gradio as gr
from fastai.vision.all import *


pasta_shapes = [
    "spaghetti",
    "penne",
    "fettuccine",
    "rigatoni",
    "bowtie",
    "linguine",
    "ravioli",
    "lasagna",
    "rotini",
]

not_pasta = [
    "rice", "quinoa", "couscous", 
    "bread", "pastries", "cookies", 
    "vegetables", "fruits", "salads", 
    "stir-fry dishes", "pizza", "sandwiches", 
    "kitchen utensils", "appliances", "cookware", 
    "furniture", "home decor", "office supplies", 
    "landscapes", "cityscapes", "forests", 
    "plants", "flowers", "trees", 
    "animals", "birds", "insects", 
    "abstract art", "patterns", "textures", 
    "fabric textures", "wood textures", "metal textures"
]



if __name__ == '__main__':
    categories = ['bowtie', 'fettuccine', 'lasagna', 'linguine', 'not pasta', 'penne', 'ravioli', 'rigatoni', 'rotini', 'spaghetti']

    def get_y(file_path):
        parent_folder = file_path.parent.name

        if parent_folder in pasta_shapes:
            return parent_folder
        else:
            return "not pasta"
        
    learn = load_learner('model.pkl')
    
    def classify_image(img):
        _, _, probs = learn.predict(img)

        return dict(zip(categories, map(float, probs)))


    image = gr.Image()
    label = gr.Label()

    intf = gr.Interface(fn=classify_image, inputs=image, outputs=label)
    intf.launch(inline=False)