minimal / app.py
brandonpollack23's picture
using fastai categories vocab
252a46c
raw
history blame contribute delete
753 Bytes
# %%
from typing import Dict, List
import gradio as gr # type: ignore
from fastai.vision.learner import load_learner, Learner # type: ignore
from PIL import Image
from numpy import shape # type: ignore
# %%
learn: Learner = load_learner("bears.pkl")
# %%
def categories() -> List[str]:
return learn.dls.vocab
def classify_image(img: Image.Image) -> Dict[str, float]:
pred, idx, probs = learn.predict(img)
print(f"Predicted the image was a {pred} bear with {probs[idx].item()*100:.02f}%")
probs = map(float, probs)
return dict(zip(categories(), probs))
# %%
image_input = gr.Image(width=224, height=224)
label = gr.Label()
iface = gr.Interface(fn=classify_image, inputs=[image_input], outputs=[label])
iface.launch()
# %%