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),
             examples=["data/1.bmp", "data/2.bmp", "data/3.bmp", "data/4.bmp", "data/5.bmp", "data/6.bmp", "data/7.jpg","data/8.jpg",
             ]).launch(share=True)