import timm
import gradio as gr
from pathlib import Path
from fastai.vision.all import *

categories = ['Tomato___Bacterial_spot', 'Tomato___Early_blight', 'Tomato___Late_blight', 'Tomato___Leaf_Mold', 'Tomato___Septoria_leaf_spot', 'Tomato___Spider_mites Two-spotted_spider_mite', 'Tomato___Target_Spot', 'Tomato___Tomato_Yellow_Leaf_Curl_Virus', 'Tomato___Tomato_mosaic_virus', 'Tomato___healthy']

learner = load_learner(Path('tomato_resnet50.pkl'))

def classify(img):
    category, index, probs = learner.predict(img)
    return (dict(zip(categories, map(float, probs))))

image = gr.inputs.Image(shape=(224))
label = gr.outputs.Label()
example = ['Tomato___Bacterial_spot.JPG', 'Tomato___Early_blight.JPG']

interface = gr.Interface(fn=classify, 
                         inputs='image', 
                         outputs='label',
                         example=example)

interface.launch()