from fastai.learner import load_learner
from fastai.vision.core import PILImage
import gradio as gr

learn = load_learner('export.pkl')

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

title = "Turtle Pleasure"
description = 'An image classifier for recognizing turtle species, created using <a href="fast.ai">FastAI</a>.'
article = """
<p>The model is a pre-trained ResNet18 fine-tuned on a dataset obtained from
<a href="https://www.inaturalist.org">iNaturalist</a> and <a href="https://www.gbif.org/">GBIF</a>.
It recognizes the following species (according to GBIF, the 50 most observed turtles in North America):</p>

<ul><li>""" + "</li><li>".join(labels) + """</li></ul>

<h3>References:</h3>
<p>GBIF.org (8 September 2023) GBIF Occurrence Download <a href="https://doi.org/10.15468/dl.mvss9v">https://doi.org/10.15468/dl.mvss9v</a></p>
"""

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(shape=(400, 400)), 
    outputs=gr.Label(num_top_classes=3),
    title=title,
    description=description,
    article=article,
    examples=['florida_softshell.jpg', 'hawksbill.jpg', 'blandings.jpg']
)
iface.launch()