import gradio as gr
from fastai.vision.all import load_learner
from PIL import Image
import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath

model = load_learner('export.pkl')

def classify_image(img):
    """Classifies an image according to three categories: dung beetle, elephant, or dolphin.

    Args:
        img (any): Any image will be converted to expected type. 

    Returns:
        _type_: Probabilies according to the three types. 
    """
    # Convert the image to a format the model expects
    img = Image.fromarray(img.astype('uint8'), 'RGB')
    # Make a prediction
    probs = model.predict(img)
    # Return the result
    return {model.dls.vocab[i]: float(probs[i]) for i in range(len(model.dls.vocab))}

demo = gr.Interface(
    title = "A dung beetle / dolphin / elephant image classifier",
    fn=classify_image, 
    inputs = gr.Image(
        label = 'Upload an image of a dung beetle, a dolphin, or an elephant!'),
    outputs="label")

if __name__ == "__main__":
    demo.launch()