pet_classifier / app.py
nkhimin
added categorization
4d52199
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()