# Load the model from gluten.py
# Create a gradio interface for the model
# Run the model on the interface

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

import fastai


# This method is required for unpickling
def label_for_path(path):
    return {
        'bread': 'glutenful',
        'gluten_free': 'glutenfree',
        'man-made': 'glutenfree',
        'fruit': 'glutenfree',
        'malt_beverage': 'glutenful',
        'meat': 'glutenfree',
        'soy_sauce': 'glutenful',
        'tamari': 'glutenfree'
    }[str(list(path.parts)[1])]


learn_inf = load_learner(Path() / 'gluten.pkl')
categories = ('gluten-free', 'glutenful')


def classify_image(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn_inf.predict(img)
    return {categories[i]: float(probs[i]) for i in range(len(categories))}


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

# examples = all files in the /examples folder
examples = [f"examples/{i}" for i in os.listdir("examples")]

intf = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    examples=examples
)

if __name__ == "__main__":
    intf.launch()