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

learn = load_learner('snakes-v2-457423.pkl')

categories = ['Banded Krait', 
              'Brahminy Blindsnake', 
              'Buff Striped Keelback', 
              'Burmese Python', 
              'Chinese Cobra', 
              'Chinese Green Snake', 
              'Chinese Green Tree Viper', 
              'Common Mock Viper', 
              'Copperhead Rat Snake', 
              'Golden Tree Snake', 
              'Indo-Chinese Rat Snake', 
              'King Cobra', 
              'Kramer’s Pit Viper', 
              'Malayan Pit Viper', 
              'Many-banded Krait', 
              'Oriental Ratsnake', 
              'Oriental Whipsnake', 
              'Painted Bronzeback', 
              'Reticulated Python', 
              'Small-banded Kukri Snake', 
              'White-lipped Pit Viper', 
              'Yellow-spotted Keelback']

def classify_snake(image):
    pred, idx, probs = learn.predict(image)
    return dict(zip(categories, map(float,probs)))

image = gr.inputs.Image(shape=(192, 192))
label = gr.outputs.Label()
examples = ['golden-tree-snake.jpg', 
            'painted-bronzeback.jpg', 
            'reticulated-python.jpg', 
            'chinese-green-tree-viper.jpg']	

intf = gr.Interface(fn=classify_snake, inputs=image, outputs=label, examples=examples)
intf.launch()