from fastai.collab import *
from fastai.tabular.all import *
import gradio as gr
from PIL import Image


# Load pretrained fastai Learner object
learn = load_learner('model.pkl', cpu=True)

# Used to convert prediction classes to text
part_labels = ['chert1-2mm','obsidian1-2mm', 'soil2-4mm']

def predict(inp):
    """
    Prediction for fast.ai fine tuned particle model
    """
    prediction = learn.predict(inp)[2]
    confidences = {part_labels[i]: float(prediction[i]) for i in range(3)}
    return confidences

gr.Interface(fn=predict, 
             inputs=gr.Image(type="numpy"),
             outputs=gr.Label(num_top_classes=3)).launch()
             #examples=["Default-archeo_transparency_EXP00006-Obsidian-HH-2mm_20210504_131710(All particles)5300_57_05.bmp"])