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

snake_labels = (
    "Monocled cobra",
    "Egyptian cobra",
    "Black-necked spitting cobra",
    "Samar cobra",
    "Red spitting cobra",
    "Javan spitting cobra",
    "Spectacled cobra",
    "Russell's viper",
    "Horned vipers",
    "Bush vipers",
    "Eyelash viper",
    "Saw-scaled vipers",
    "Banded krait",
    "Black mamba",
    "Inland taipan",
    "Eastern brown snake",
    "Rattle snake",
    "King cobra"
)

model = load_learner('snake-recognizer-v0.pkl')

def recognize_snake(photo):
    pred,idx, probs = model.predict(photo)
    return pred




image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label(num_top_classes=5)
examples = [
    'viper-2.jpg',
    'shutterstock_2062214282-edited-1-scaled.jpg',
    'download (3).jpg',
    'download (4).jpg',
    'Naja_sputatrix.jpg',
    'download (6).jpg',
    'download.jpg',
    'Naja-pallida-by-Wikimedia-commons.jpg',
    'istockphoto-1358737491-612x612.jpg',
    'istockphoto-916895080-612x612.jpg',
    'images.jpg',
    'high.jpg',
    'download (2).jpg',
    'download (5).jpg',
    'black-necked-spitting-cobra-naja-nigricollis-wklein.jpg',
    '4538a8531287f0a3ab464d0b9dd69744.jpg',
    'Bothriechis_schlegelii_(La_Selva_Biological_Station).jpg'
    ]

iface = gr.Interface(fn=recognize_snake, inputs=image, outputs=label, examples=examples)
iface.launch(inline=False)