|
|
import gradio as gr |
|
|
import pickle |
|
|
from PIL import Image |
|
|
|
|
|
from fastai.learner import load_learner |
|
|
|
|
|
model = load_learner('./mushrooms.pkl') |
|
|
|
|
|
|
|
|
categories = [ |
|
|
"Agaricus", |
|
|
"Amanita", |
|
|
"Boletus", |
|
|
"Cortinarius", |
|
|
"Entoloma", |
|
|
"Hygrocybe", |
|
|
"Lactarius", |
|
|
"Russula", |
|
|
"Suillus", |
|
|
] |
|
|
|
|
|
def classify_image(image): |
|
|
prediction, _, probs = model.predict(image) |
|
|
return dict(zip(categories, map(float, probs))) |
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=classify_image, |
|
|
inputs=gr.Image(type="pil"), |
|
|
outputs=gr.Label(), |
|
|
title="Image Classifier", |
|
|
description="WHAT IS THE MUSHROOM?" |
|
|
) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
iface.launch() |
|
|
|
