fmussari's picture
Shared in the forum
d4c65cf
# AUTOGENERATED! DO NOT EDIT! File to edit: . (unless otherwise specified).
__all__ = ['learn', 'classify_image', 'categories', 'title', 'description', 'image', 'label', 'jpg_files', 'examples',
'intf']
# Cell
from fastai.vision.all import *
import gradio as gr
# Cell
learn = load_learner('tower_parts_model.pkl')
# Cell
categories = learn.dls.vocab
def classify_image(img):
pred, idx, probs = learn.predict(img)
return dict(zip(categories, map(float,probs)))
# Cell
title = "Classify Telecommunication Tower Parts"
description = "This deep learning model was trained with fastai using only 478 images of telecommunication towers parts."
description += "\nThe model was trained to recognize the following 8 categories:"
description += "\n- Base plate | Grounding bar | Identification | Ladder | Light | Lightning rod | Platform | Transmission lines"
description += "\n- You can test the model with the given examples (see below), or upload your own pictures."
# Cell
image = gr.inputs.Image(
shape=(200,200),
label='Load image'
)
label = gr.outputs.Label()
#examples = ['DSC01955.jpg', 'DSC01956.jpg']
jpg_files = os.listdir()
examples = [file for file in jpg_files if file.endswith('.jpg')]
# Cell
intf = gr.Interface(
fn=classify_image,
inputs=image,
outputs=label,
examples=examples,
title=title,
description=description,
)
intf.launch(inline=False)