import gradio as gr
import pandas as pd
# import pathlib
from fastbook import *
from fastai.vision.widgets import *
from fastai.vision.all import *


# Load the model. "PosixPath" is something with windows/linux, I dont know really.
# temp = pathlib.PosixPath
# pathlib.PosixPath = pathlib.WindowsPath
learn = load_learner('model.pkl') # Load the model itself
# pathlib.PosixPath = temp
categories = learn.dls.vocab # Get the list of labels from the model

# Load the list of trash-sorting, use the items as index.
df_sort = pd.read_csv('Lista.csv',sep =";").set_index('Avfall')

def classify_image(img):
    # Make the prediction
    trash,idx,probs = learn.predict(PILImage.create(img)) # Make prediction
    df = pd.DataFrame()    # Create dataframe
    df['categories'] = categories # Add categories to dataframe
    df['probabilities'] = probs.numpy() # Add probabilities to dataframe
    sorted_df = df.sort_values(by=['probabilities'], ascending=False).head() # Sort by probability, highest first, take the top 5
    predictions = dict(zip(sorted_df['categories'].tolist(),map(float,sorted_df['probabilities'].tolist()))) # Now convert to a dictionary that we return later
    
    
    # Create sorting statement
    sort_text = "Sorteras som " + df_sort.loc[trash].tolist()[0]
    return "Det där är...", predictions, sort_text # Return the dictionary


image = gr.Image(type='pil')
label = ["text",gr.Label(),"text"]

iface = gr.Interface(fn=classify_image, inputs=image, outputs=label)
iface.launch()