import gradio as gr
from fastai.vision.all import *

learn = load_learner("resnet18.pkl")

categories = ("Healthy", "Peacock Spot")


def classify_health(input_img):
    pred, idx, probs = learn.predict(input_img)
    return dict(zip(categories, map(float, probs)))


labels = gr.Label()
examples = [
    "examples/healthy.jpg",
    "examples/healthy2.jpg",
    "examples/peacock_spot.jpeg",
]
demo = gr.Interface(
    classify_health,
    inputs=gr.Image(height=224, width=224),
    outputs=labels,
    examples=examples,
)
demo.launch(inline=False)