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-v3.pkl"

model = load_learner(model_dir)

parrot_species = ['african grey parrot',
 'australian king parrot',
 'blue lorikeet',
 'blue-and-yellow macaw',
 'blue-headed parrot',
 'budgerigar',
 'burrowing parrot',
 'caique parrot',
 'catalina macaw',
 'chestnut-fronted macaw',
 'cockatiels',
 'crimson rosella',
 'cuban amazon',
 'eclectus parrot',
 'galah',
 'golden parakeet',
 'great green macaw',
 'great hanging parrot',
 'greater vasa parrot',
 'hahn_s macaws',
 'hooded parrot',
 'hyacinth macaw',
 'kea',
 'kākāpō',
 'lovebirds',
 'monk parakeet',
 'orange-winged amazon',
 'palm cockatoo',
 'parrotlet',
 'plum-headed parakeet',
 'puerto rican amazon',
 'rainbow lorikeet',
 'red-breasted parakeet',
 'red-crowned amazon',
 'red-crowned parakeet',
 'red-fan parrot',
 'red-shouldered macaw',
 'red-tailed black cockatoos',
 'rose-ringed parakeet',
 'saint vincent amazon',
 'scarlet macaw',
 'senegal parrot',
 'spixs macaw',
 'sun conure',
 'thick-billed parrot',
 'turquoise-fronted amazon',
 'vernal hanging parrot',
 'white cockatoo',
 '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",
    ]

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