import gradio as gr
from pathlib import Path
# from fastai.vision.all import *   # noqa: F403
from fastai.learner import load_learner
from fastai.vision.core import PILImage
import os


# import skimage


# Define any custom functions or classes that the model depends on
def is_cat(x):  # Make sure to define this correctly as it was used during training
    return x[0].isupper()

print(os.path.abspath("model-export.pkl"))
learn = load_learner("model-export.pkl")

labels = learn.dls.vocab

categories = ('Dog', 'Cat')
def classify_image(img):
    img = PILImage.create(img)
    pred,idx,probs = learn.predict(img)
    return dict(zip(categories, map(float,probs)))


title = "Pet Breed Classifier"
description = "A pet breed classifier trained on the Oxford Pets dataset with fastai. Created as a demo for Gradio and HuggingFace Spaces."
article = "<p style='text-align: center'><a href='https://tmabraham.github.io/blog/gradio_hf_spaces_tutorial' target='_blank'>Blog post</a></p>"
path = Path('examples')
allowed_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.gif'}
examples = [file for file in path.iterdir() if file.suffix.lower() in allowed_extensions]

gr.Interface(
    fn=classify_image,
    inputs="image",
    outputs="label",
    title=title,
    description=description,
    article=article,
    examples=examples
).launch()