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

learn = load_learner('curlyhairmodelv3.pkl')

categories = ('Curly', 'Straight')

def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    return dict(zip(categories, map(float,probs)))

title = "Hair Type Classifier"
description = "A hair type classifier trained on the Oxford Pets dataset with fastai."
examples = [['c1.jpeg'],['c2.jpeg'],['c3.jpeg'],
            ['c4.jpeg'],['c5.jpeg'],['c6.jpeg'],
            ['c7.jpeg'],['c9.jpeg'],['c10.jpeg'],
            ['c11.jpeg'],['c12.jpeg'],
            ['s2.jpeg'],['s4.jpeg'],['s5.jpeg'],
            ['s6.jpeg'],['s7.jpeg'],['s12.jpeg'],
            ['sc1.jpeg'],['sc2.jpeg'],['sc3.jpeg'],
            ['sc4.jpeg'],['sc5.jpeg'],['sc6.jpeg']
			 ]
interpretation='default'
enable_queue=True

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