drifterz13's picture
Set up project
06f05e4
import gradio as gr
from fastai.vision.all import *
import os
learn = load_learner('model.pkl')
labels = learn.dls.vocab
def predict(img):
img = PILImage.create(img)
pred, pred_idx, probs = learn.predict(img)
return {labels[i]: float(probs[i]) for i in range(len(labels))}
ui = gr.Interface(
inputs=gr.Image(type="pil", label="Upload Image", width=512, height=512),
fn=predict,
outputs=gr.Label(num_top_classes=3),
flagging_options=["blurry", "incorrect", "other"],
examples=[
os.path.join(os.path.dirname(__file__), "samples", "pikachu.png"),
os.path.join(os.path.dirname(__file__), "samples", "raichu.jpeg"),
os.path.join(os.path.dirname(__file__), "samples", "mimikyu.jpeg"),
os.path.join(os.path.dirname(__file__), "samples", "charmander.jpeg"),
os.path.join(os.path.dirname(__file__), "samples", "bulbasaur.png"),
os.path.join(os.path.dirname(__file__), "samples", "squirtle.png"),
os.path.join(os.path.dirname(__file__), "samples", "dunno.jpeg"),
]
)
ui.title = "Pikachu or Raichu or Mimikyu?"
ui.description = "This is a simple image classifier that can classify between Pikachu, Raichu and Mimikyu."
ui.launch(share=True)