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

learn = load_learner('learner.pkl')
labels = learn.dls.vocab

def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

for root, dirs, files in os.walk(r'sample_images/'):
    for filename in files:
        print(filename)

title = "Paddy Disease Classifier with EdgeNeXt"
description = "9 Diseases + 1 Normal class."
interpretation='default'
examples = ["sample_images/"+file for file in files] 
article="<p style='text-align: center'><a href='https://dicksonneoh.com/portfolio/pytorch_at_the_edge_timm_torchscript_flutter/' target='_blank'>来源博客</a></p>"
enable_queue=True

gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(224, 224)),
    outputs=gr.outputs.Label(num_top_classes=5),
    title=title,
    description=description,
    article=article,
    examples=examples,
    interpretation=interpretation,
    enable_queue=enable_queue,
    theme="grass",
).launch()