import gradio as gr
from fastai.vision.all import *

# Define the main prediction function:
learn = load_learner('./model_params/planet_cls_1.pkl')
labels = learn.dls.vocab
def planet_classifier_predict(img) -> str:
  img = PILImage.create(img)
  _, _, probs = learn.predict(img)
  return {labels[i]: float(probs[i]) for i in range(len(labels))}


if __name__ == "__main__":
    # General web app params:
    title = "Planet Classifier Demo"
    description = "Let's say you want to classify a picture if it is one of the 9 planets: \
                   Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune. \
                   You can do it using this awesome tool!!!! Try it :)"
    article = "<p style='text-align: center'><a href='https://www.tanishq.ai/blog/posts/2021-11-16-gradio-huggingface.html' target='_blank'>The main source of inspiration :) </a></p>"
    examples = ['./example_images/saturn.jpeg', './example_images/car.jpeg']
    interpretation = 'default'
    enable_queue = True
    inputs = "image"
    outputs = "label"
    # The main app interface:
    gr.Interface(
        fn=planet_classifier_predict,
        inputs=inputs,
        outputs=outputs,
        title=title,
        description=description,
        article=article,
        examples=examples,
    ).launch()