import gradio as gr
import torch
from huggingface_hub import from_pretrained_fastai
from pathlib import Path

examples = ["image_1.png", "image_2.png", "image_3.png", "image_4.png", "image_5.png"]
repo_id = "hugginglearners/grapevine_leaves_classification"
path = Path("./")

def get_y(r):
    return r["label"]
    
def get_x(r):
    return path/r["fname"]
    
learner = from_pretrained_fastai(repo_id)

def inference(image):
    label_predict,_,probs = learner.predict(image)
    return f"This grapevine leave is {label_predict} with {100*probs[torch.argmax(probs)].item():.2f}% probability"

gr.Interface(
    fn=inference,
    title="Grapevine leave image classification",
    description = "Predict which type of grapevine leave belong to Ak, Ala_Idris, Buzgulu, Dimnit, Nazli",
    inputs="image",
    examples=examples,
    outputs=gr.Textbox(label='Prediction'),
    cache_examples=False,
    article = "Author: <a href=\"https://huggingface.co/vumichien\">Vu Minh Chien</a>",
).launch(debug=True, enable_queue=True)