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

learn = load_learner('export_ninjago_classifier.pkl')

labels = ['Acidicus_Serpentine', 'Acronix_Vermillion', 'Akita_Never-Realm inhabitants', 'Antonia_Recurring allies', 'Ash_Elemental Masters', 'Aspheera_Pyro Vipers', 'Bansha_Ghosts', 'Benthomaar_Merlopians', 'Blunck_Vermillion', 'Bolobo_Elemental Masters', 'Bucko_Sky Pirates', 'Captain Soto_Citizens of Ninjago City', 'Chamille_Elemental Masters', 'Char_Pyro Vipers', 'Chew Toy_Dragon Hunters', 'Chief Mammatus_The Island inhabitants', 'Chope_Anacondrai Cultists', 'Clancee_Sky Pirates', 'Clouse_Anacondrai Cultists', 'Clutch Powers_Recurring allies', 'Cole_Main Characters', 'Cyren_Sky Pirates', 'Cyrus Borg_Recurring allies', 'Daddy No Legs_Dragon Hunters', 'Dareth_Recurring allies', 'Dogshank_Sky Pirates', 'Doubloon_Sky Pirates', 'Dr. Julien_Family members', 'Edna Walker_Family members', 'Ed Walker_Family members', 'Eyezor_Anacondrai Cultists', 'Faith_Dragon Hunters', 'Fangtom_Serpentine', 'Flintlocke_Sky Pirates', 'Fred Finely_Citizens of Ninjago City', 'Fuchsia Ninja_New Ninja', 'Fugi-Dove_Citizens of Ninjago City', 'General Arcturus_Serpentine', 'General Cryptor_Nindroids', 'General Kozu_Stone Army', 'General Vex_Never-Realm inhabitants', 'Ghoultar_Ghosts', 'Gleck_Shintaro Kingdom inhabitants', 'Glutinous_Merlopians', 'Gravis _Elemental Masters', 'Griffin Turner_Elemental Masters', 'Hounddog McBrag_Citizens of Ninjago City', 'Jacob Pevsner_Elemental Masters', 'Jay_Main Characters', 'Jet Jack_Dragon Hunters', 'Kai_Main Characters', 'Kapau_Anacondrai Cultists', 'Karlof_Elemental Masters', 'Khanjikhan_Sky Pirates', 'Killow_Sons of Garmadon', 'King Trimaar_Merlopians', 'King Vangelis_The Skull Sorcerer_Shintaro Kingdom inhabitants', 'Kruncha_Skulkin', 'Krux_Vermillion', 'lables.txt', 'Lloyd Garmadon_Main Characters', 'Lord Garmadon_Main Characters', 'Luke Cunningham_Sons of Garmadon', 'Machia_Vermillion', 'Master Chen_Anacondrai Cultists', 'Master Wu_Main Characters', 'Maya_Family members', 'Milton Dyer_Prime Empire inhabitants', 'Min-Droid_Nindroids', 'Misako_Family members', 'Mistaké_Recurring allies', 'Mister F_Council of the Crystal King', 'Monkey Wretch_Sky Pirates', 'Morro_Ghosts', 'Mr. E_Sons of Garmadon', 'Murt_Shintaro Kingdom inhabitants', 'Muzzle_Dragon Hunters', 'Nadakhan_Sky Pirates', 'Nelson_Recurring allies', 'Neuro_Elemental Masters', 'Nuckal_Skulkin', 'Nya_Main Characters', 'Officer Toque_Recurring allies', 'Okino_Prime Empire inhabitants', 'Orange Ninja_New Ninja', 'Paleman_Elemental Masters', 'Patty Keys_Citizens of Ninjago City', 'Pink Ninja_New Ninja', 'P.I.X.A.L._Main Characters', 'Police Commissioner_Recurring allies', 'Postman_Citizens of Ninjago City', 'Prince Kalmaar_Merlopians', 'Princess Harumi_Sons of Garmadon', 'Princess Vania_Shintaro Kingdom inhabitants', 'Pythor P. Chumsworth_Serpentine', 'Raggmunk_Vermillion', 'Ray_Family members', 'Ronin_Recurring allies', 'Rufus MacAllister_Citizens of Ninjago City', 'Samukai_Skulkin', 'Scott_Prime Empire inhabitants', 'Sensei_Master Kodokuna Yang_Ghosts', 'Shade_Elemental Masters', 'Skales_Serpentine', 'Skalidor_Serpentine', 'Skylor Chen_Elemental Masters', 'Slithraa_Serpentine', 'Soul Archer_Ghosts', 'Squiffy_Sky Pirates', 'Sushimi_Prime Empire inhabitants', 'The First Spinjitzu Master_Family members', 'The Great Devourer_Serpentine', 'The Ice Emperor_Never-Realm inhabitants', 'The Iron Baron_Dragon Hunters', 'The Mechanic_Prime Empire inhabitants', 'The Omega_Oni', 'The Overlord_Stone Army', 'The Preeminent_Ghosts', 'Timothy "Twitchy Tim" Batterson_The Island inhabitants', 'Tox_Elemental Masters', 'Ultra Violet_Sons of Garmadon', 'Unagami_Prime Empire inhabitants', 'Wrayth_Ghosts', 'Wyplash_Skulkin', 'Zane_Main Characters', 'Zugu_Anacondrai Cultists']




def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}


title = "Ninjago Character Classifier"
description = "Guesses the name of the Ninjago characters. Created from the fastai 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 = ['Lloyd.jpg', 'Cole.png']
interpretation = 'default'
enable_queue = True

gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(192, 192)), outputs=gr.outputs.Label(num_top_classes=3), title=title,
             description=description, examples=examples, interpretation=interpretation, enable_queue=enable_queue).launch()