Redhotchilipoppy's picture
Added fastbook back in
c06002b verified
raw
history blame
1.51 kB
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()