import gradio as gr
from fastai.vision.all import load_learner
from PIL import Image

from contextlib import contextmanager
import pathlib

@contextmanager
def set_posix_posix():
    windows_backup = pathlib.WindowsPath
    try:
        pathlib.WindowsPath = pathlib.PosixPath
        yield
    finally:
        pathlib.PosixPath = windows_backup

EXPORT_PATH = pathlib.Path("export.pkl")

with set_posix_posix():
    learn_inf = load_learner(EXPORT_PATH)

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
    pred_class, pred_idx, probs = learn_inf.predict(img)
    # Return the result
    return {learn_inf.dls.vocab[i]: float(probs[i]) for i in range(len(learn_inf.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()