import os

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


learner = load_learner(hf_hub_download("pinkpekoe/lesson2-bear-classifier", "export.pkl"))

def predict_bear_type(img_path):
    img = PILImage.create(img_path)
    pred, pred_idx, probs = learner.predict(img)
    probabilities = [
        f"{probs[i]:.02f}*" if i == pred_idx else f"{probs[i]:.02f}"
        for i in range(len(probs))
    ]
    return f"Prediction: {pred}; Probabilities: " + ", ".join(probabilities)


title = "Pet Breed Classifier"
description = "A pet breed classifier trained on the Oxford Pets dataset with fastai. Created as a demo for Gradio and HuggingFace Spaces."
article = "<p style='text-align: center'><a href='https://tmabraham.github.io/blog/gradio_hf_spaces_tutorial' target='_blank'>Blog post</a></p>"
examples = ["black.jpeg", "brown.jpeg", "teddy.jpeg"]
interpretation = "default"
enable_queue = True

iface = gr.Interface(
    fn=predict_bear_type,
    inputs="image",
    outputs="text",
    title=title,
    description=description,
    article=article,
    examples=examples,
    interpretation=interpretation,
)
iface.launch(enable_queue=enable_queue)