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

# Download the model from Hugging Face
model_path = hf_hub_download(repo_id="principle/bears-classifier-model", filename="bears_modle.pkl")

# Load the model
learn = load_learner(model_path)

# Define the prediction function
def classify_bear(img):
    pred, pred_idx, probs = learn.predict(img)
    return {
        'black': float(probs[0]),
        'grizzly': float(probs[1]),
        'teddy': float(probs[2])
    }

# Create the Gradio interface
iface = gr.Interface(
    fn=classify_bear,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=3),
    title="Bear Classifier",
    description="Upload an image to classify the type of bear: grizzly, black, or teddy."
)
# !
# Launch the app
iface.launch()