Spaces:
Sleeping
Sleeping
from fastai.vision.all import load_learner | |
import gradio as gr | |
cap_labels = cap_labels = { | |
'baseball cap', | |
'beanie cap', | |
'fedora cap', | |
'cowboy hat', | |
'kepi cap', | |
'flat cap', | |
'trucker cap', | |
# 'newsboy cap' | |
'pork pie hat', | |
'bowler hat', | |
'top hat', | |
'sun hat', | |
'boater hat', | |
# 'ivy cap', | |
'bucket hat', | |
'balaclava cap', | |
'turban cap', | |
'taqiyah cap', | |
'rasta cap', | |
'visor cap' | |
} | |
version = 1 | |
model_path = f"cap-recognizer-v{version}.pkl" | |
model = load_learner(model_path) | |
def recognize_image(image): | |
pred, idx, probs = model.predict(image) | |
return dict(zip(sorted(cap_labels), map(float, probs))) | |
image = gr.inputs.Image(shape=(192, 192)) | |
label = gr.outputs.Label() | |
examples = [ | |
'test_images/test_0.jpg', | |
'test_images/test_1.jpg', | |
'test_images/test_2.jpg', | |
'test_images/test_3.jpg', | |
'test_images/test_4.jpg', | |
'test_images/test_5.jpg'] | |
iface = gr.Interface(fn=recognize_image, inputs=image, outputs=label, examples=examples) | |
iface.launch(inline=False) |