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"

learn = from_pretrained_fastai(repo_id)

categories = learn.dls.vocab

def predict(image):
    label, index, probs = learn.predict(image)
   
    return dict(zip(categories, map(float,probs)))

title = "Snail, snake, slug Classifier"
description = f"""<h1> Slug, snake, snail or other<h1>
              <p>A classifier trained on about 600 images. Created as a demo for Gradio and HuggingFace Spaces."""
examples = ['330px-Orange_slug.jpg', 'Green_Snakes.jpg', 'Helix_pomatia_002.JPG', 'baum-jung.jpg']



interpretation='default'
enable_queue=True


gr.Interface(
    predict,
    inputs=gr.components.Image(label="candidate", type="filepath"),
    outputs=gr.components.Label(num_top_classes=4),
    title=title,
    examples=examples,
    description=description,
    # interpretation=interpretation,
    # enable_queue=enable_queue
).launch()