28 lines
872 B
Python
28 lines
872 B
Python
|
import gradio as gr
|
||
|
|
import os
|
||
|
|
|
||
|
|
from fastai.learner import load_learner
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
path = Path(os.path.dirname(os.path.realpath(__file__)))
|
||
|
|
|
||
|
|
# Make sure to export the model first via the notebook.
|
||
|
|
# The exported model is not included in the repository.
|
||
|
|
model = load_learner(path/'export.pkl')
|
||
|
|
labels = model.dls.vocab
|
||
|
|
|
||
|
|
def predict(img):
|
||
|
|
_, _, probabilities = model.predict(img)
|
||
|
|
return {labels[i]: float(probabilities[i]) for i in range(len(labels))}
|
||
|
|
|
||
|
|
title = "Grizzly, Black, or Teddy Bear?"
|
||
|
|
description = "This is a simple image classifier that can predict whether a given image contains a grizzly bear, black bear, or teddy bear."
|
||
|
|
|
||
|
|
gr.Interface(
|
||
|
|
fn = predict,
|
||
|
|
inputs = gr.Image(width = 512, height = 512),
|
||
|
|
outputs = gr.Label(num_top_classes = 3),
|
||
|
|
title = title,
|
||
|
|
description = description,
|
||
|
|
allow_flagging = 'never'
|
||
|
|
).launch()
|