|
|
|
|
|
|
|
__all__ = ['cat_dog_model', 'bear_model', 'input_model', 'input_image', 'output_label', 'bear_examples', 'cat_examples', |
|
'examples', 'intf', 'get_model', 'classify_image'] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from fastai.vision.all import * |
|
import gradio as gr |
|
|
|
|
|
cat_dog_model = load_learner('cat_dog.pkl') |
|
bear_model = load_learner("bear_classifier.pkl") |
|
|
|
|
|
|
|
def get_model(model_name): |
|
if model_name == 'Cat vs Dog Model': |
|
return cat_dog_model |
|
elif model_name == 'Bear Model': |
|
return bear_model |
|
else: |
|
raise ValueError("Model not found") |
|
|
|
|
|
def classify_image(model_name, img): |
|
model = get_model(model_name) |
|
categories = model.dls.vocab |
|
pred, idx, probs = model.predict(img) |
|
return dict(zip(categories, map(float, probs))) |
|
|
|
|
|
|
|
input_model = gr.Dropdown(['Cat vs Dog Model', 'Bear Model']) |
|
input_image = gr.components.Image(width=192, height=192) |
|
|
|
|
|
output_label = gr.components.Label() |
|
|
|
|
|
bear_examples = ['black.jpg', 'grizzly.jpg', 'teddy.jpg'] |
|
cat_examples = ['cat.jpg', 'dog.jpg'] |
|
|
|
examples = [ |
|
['Bear Model',bear_examples[0]], |
|
['Bear Model',bear_examples[1]], |
|
['Bear Model',bear_examples[2]], |
|
['Cat vs Dog Model', cat_examples[0]], |
|
['Cat vs Dog Model', cat_examples[1]], |
|
] |
|
|
|
|
|
intf = gr.Interface(fn=classify_image, |
|
inputs=[ |
|
input_model, |
|
input_image], |
|
outputs=output_label, |
|
examples=examples, |
|
theme=gr.themes.Ocean()) |
|
|
|
|
|
intf.launch(inline=False, |
|
|
|
) |
|
|