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

# Importing the model
model = load_learner('model.pkl')

categories = model.dls.vocab

def predict_snake_type(img):
    snake_type, idx, probs = model.predict(img)
    return dict(zip(categories, map(float, probs)))

# Gradio Interface
input_img = gr.Image()
outputs = gr.Label(num_top_classes=5)

gr.Interface(
    fn=predict_snake_type, 
    inputs=[input_img], outputs=outputs,
    examples=[
        'examples/agkistrodon-contortrix.jpg', 
        'examples/haldea-striatula.jpg', 
        'examples/masticophis-flagellum.jpg',
        'examples/storeria-occipitomaculata.jpg'
        ],
    title='Classifying Different Breeds of snakes',
    description='This is a multi-classification model that identifies different breeds of snakes.\nPlease note that is model is only `74%` accurate.'
).launch()