asquirous's picture
Fix probability of "Not Sure/Others"
ef69052
raw
history blame contribute delete
827 Bytes
from fastai.vision.all import *
import torch
import gradio as gr
learn = load_learner("model_tvdesktop.pkl")
labels = learn.dls.vocab
def classify_image(img):
img = PILImage.create(img)
pred, idx, probs = learn.predict(img)
output = dict(zip(labels, map(float, probs)))
print(probs[0].item()*100)
count = 0
for out in output:
val = output[out]
if val*100 < 60:
count += 1
if count == 2:
return {"Not sure/Others": 0}
return output
image = gr.inputs.Image(shape=(224, 224))
label = gr.outputs.Label()
title = "CRT TV and Desktop Monitor Classifier"
description = "A simple image classifier."
intf = gr.Interface(
fn=classify_image,
inputs=image,
outputs=label,
title=title,
description=description
)
intf.launch(inline=False)