Spaces:
Runtime error
Runtime error
File size: 1,338 Bytes
51d7046 e8c97c4 f39ff71 51d7046 89e0eb8 407a4df 2b21ef9 67edaec e8c97c4 0df3b7e fdcd337 89e0eb8 407a4df 67edaec e8c97c4 0df3b7e fdcd337 67edaec e8c97c4 67edaec b403237 67edaec 1cd0cab 407a4df f39ff71 407a4df cb8326d 67edaec b403237 67edaec 407a4df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import gradio as gr
from gradio.mix import Parallel
from fastai.vision.all import load_learner
def classify_image_color(img):
from fastai.vision.all import load_learner
learn = load_learner('model-color.pkl')
categories = learn.dls.vocab
pred, idx, probs = learn.predict(img)
return {f"{category}": float(prob) for category, prob in zip(categories, probs)}
def classify_image_shape(img):
from fastai.vision.all import load_learner
learn = load_learner('bricks-model.pkl')
categories = learn.dls.vocab
pred, idx, probs = learn.predict(img)
return {f"{category}": float(prob) for category, prob in zip(categories, probs)}
def classify_image(img):
color_result = classify_image_color(img)
shape_result = classify_image_shape(img)
result = {}
for key in set(color_result.keys()) | set(shape_result.keys()):
result[key] = color_result.get(key, 0.0) + shape_result.get(key, 0.0)
return result
def postprocess(prediction):
sorted_pred = sorted(prediction.items(), key=lambda x: x[1], reverse=True)
return sorted_pred
image = gr.inputs.Image(shape=(256, 256))
label = gr.outputs.Label()
intf = gr.Interface(
fn=classify_image,
inputs=image,
outputs=label,
examples="",
title="Lego Brick Classifier",
layout="vertical"
)
intf.launch()
|