import gradio as gr
import numpy as np
from huggingface_hub import from_pretrained_fastai

learn_imp = from_pretrained_fastai('MikeTrizna/Prunus_lineage_classifier')

def predict(image_np, description=None):
  classes = ['neotropical racemose', 'paleotropical racemose',
             'temperate diploid', 'temperate racemose']
  pred,pred_idx,probs = learn_imp.predict(image_np)
  confidences = {classes[idx]: f'{prob:.04f}' for idx, prob in enumerate(probs)}
  return confidences

article = """Let's include a *table* here with links to the specific samples in the examples."""

example_list = ["examples/myrtifolia_1c086296-6d1f-4218-a18a-ca2f86c295d0.jpg",
                "examples/oleifolia_104f44c9-63f1-4579-93c2-54c6ddeddeda.jpg",
                "examples/oocarpa_e3276fdd-ce58-40dc-bf20-8742f7634428.jpg",
                "examples/pullei_e5ac59df-5b18-4c8d-a19b-02131d358855.jpg",
                "examples/serotina_0060a57d-e779-4984-913c-95b576daf0d3.jpg",
                "examples/grayana_b9fa9ff9-2e8b-42f5-8728-5b0b5be98490.jpg",
                "examples/glandulosa_52e06de9-c035-4afc-9fac-c8bee7628d38.jpg",
                "examples/pensylvanica_0b13e5eb-73da-4e0e-9b1a-794b737f7716.jpg"]

demo = gr.Interface(predict,
                    gr.Image(type="numpy"),
                    outputs="label",
                    examples=example_list,
                    description="This model can place herbarium sheet images from the genus Prunus into one of four categories that represent key clades within the genus: Solitary/Corymbose, Temperate Racemose, Neotropical Racemose, Paleotropical Racemose.",
                    title="Image classification of four major lineages in the plum genus",
                    article=article)

demo.launch()