import gradio as gr
from transformers import pipeline

from huggingface_hub import from_pretrained_fastai

# repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
repo_id = "MasleK/snails_snakes_slugs"

learner = from_pretrained_fastai(repo_id)

def predict(image):
    label, index, scores = learner.predict(image)
   
    return {l: scores[i].item() for i,l in enumerate(learner.dls.vocab)}
title = "Snail, snake, slug Classifier"
description = "A  classifier trained on about 300 images. Created as a demo for Gradio and HuggingFace Spaces."
examples = ['330px-Orange_slug.jpg', 'Green_Snakes.jpg', 'Helix_pomatia_002.JPG']


gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="candidate", type="filepath"),
    outputs=gr.outputs.Label(num_top_classes=3),
    title=title,
    examples=examples,
    description=description,
).launch()