__all__ = ['get_y','learn', 'classify_images', 'image', 'label',
'intf']

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

def get_y(o):
  return [o.parent.name]

learn = load_learner('multi_label_owl.pkl')

# categories = ('barn', 'barred', 'snowy')

def classify_images(img):
    pred,idx,probs = learn.predict(img)
    idxs = torch.where(probs>=0.90)
    if len(idxs[0])==0:
        return 'None'
    elif (len(idxs[0]>1)):
        prob = sorted(probs, reverse=True)
        idxs = torch.where(probs==prob[0])
        return "cat: "+learn.dls.vocab[idxs][0]+" ; "+"prob: "+str(probs[idxs].item())
    else:
        return "cat: "+learn.dls.vocab[idxs][0]+" ; "+"prob: "+str(probs[idxs].item())

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

intf = gr.Interface(title="Owl Detector(Multilabel)",
            description = "An image classifier that detects three types of owl(snowy, barn, barred) but have been trained using multilabel classification",
            fn=classify_images, 
            inputs=image, 
            outputs=label)
intf.launch(inline=False)