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

wandb.login(key='4d2ff70f59f0bfc0397e9fea4300819244bf3fce')
wandb.init(
    name='pepper_fastai_gradio',
    project='plant_disease_detection',
    entity='jumashafara',
)

categories = ['bacterial_spot', 'healthy']

learner = load_learner(Path('pepper_fastai_model.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 = []

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

interface.launch()