from fastai.vision.all import *
from fastai.vision.all import load_learner
import fastai
import os
import gradio as gr
import pathlib


# from google.colab import drive
# drive.mount('/content/drive/')

temp = pathlib.WindowsPath
pathlib.WindowsPath = pathlib.PosixPath

model_dir = "models/parrot-recognizer-v10.pkl"

model = load_learner(model_dir)

parrot_species = ['african grey parrot',
 'australian king parrot',
 'australian night parrot',
 'bare eyed cockatoo',
 'blue and yellow macaw',
 'blue headed parrot',
 'blue lorikeet',
 'brown hooded parrot',
 'budgerigar',
 'burrowing parrot',
 'caique parrot',
 'catalina macaw',
 'chestnut-fronted macaw',
 'citron cockatoo',
 'cockatiels',
 'crimson rosella',
 'cuban amazon',
 'eclectus parrot',
 'galah cockatoo',
 'gang gang cockatoo',
 'golden parakeet',
 'great green macaw',
 'great hanging parrot',
 'greater vasa parrot',
 'hahn_s macaws',
 'hooded parrot',
 'hyacinth macaw',
 'kea',
 'kākāpō',
 'lovebirds',
 'major mitchell_s cockatoo',
 'monk parakeet',
 'musk lorikeet',
 'palm cockatoo',
 'parrotlet',
 'plum headed parakeet',
 'puerto rican amazon',
 'rainbow lorikeet',
 'red breasted parakeet',
 'red crowned amazon',
 'red crowned parakeet',
 'red fan parrot',
 'red lory',
 'red rumped parrot',
 'red shouldered macaw',
 'red tailed black cockatoos',
 'rose ringed parakeet',
 'saint vincent amazon',
 'salmon crested cockatoo',
 'scarlet macaw',
 'senegal parrot',
 'spixs macaw',
 'sulpher crested cockatoo',
 'sun conure',
 'thick billed parrot',
 'turquoise fronted amazon',
 'umbrella cockatoo',
 'vernal hanging parrot',
 'yellow collared macaws',
 'yellow headed amazon']

def recognize_image(image):
  pred, idx, probs = model.predict(image)
  return dict(zip(parrot_species, map(float, probs)))


# im =  "/content/drive/MyDrive/Learnings/fai/test_images/unknown_12.jpg"
# img = PILImage.create(im)
# img.thumbnail((192,192))
# img

# recognize_image(img)



image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label(num_top_classes=5)

examples = [
    "test_images/unknown_00.jpg",
    "test_images/unknown_01.jpg",
    "test_images/unknown_02.jpg",
    "test_images/unknown_03.jpg",
    "test_images/unknown_04.jpg",
    "test_images/unknown_05.jpg",
    "test_images/unknown_06.jpg",
    "test_images/unknown_07.jpg",
    "test_images/unknown_08.jpg",
    "test_images/unknown_09.jpg",
    "test_images/unknown_10.jpg",
    "test_images/unknown_11.jpg",
    "test_images/unknown_12.jpg",
    "test_images/unknown_13.jpg",
    "test_images/unknown_14.jpg",
    "test_images/unknown_15.jpg",
    "test_images/unknown_16.jpg",
    "test_images/unknown_17.jpg",
    "test_images/unknown_18.jpg",
    "test_images/unknown_19.jpg",
    ]

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