Spaces:
Sleeping
Sleeping
import gradio as gr | |
import os | |
import torch | |
import timeit | |
from model import create_effnetb2_model | |
from timeit import default_timer as timer | |
from typing import Tuple,Dict | |
#setup classnames | |
class_names=['pizza','steak','sushi'] | |
#model and transforms preparation | |
effnetb2,eff_b2_trans=create_effnetb2_model(num_classes=3,seed=42) | |
#load saved weights | |
effnetb2.load_state_dict(torch.load(f="eff_model_food24.pth", | |
map_location=torch.device('cpu'))) | |
#write predict function | |
def predict(img): | |
#start a timer | |
start_timer=timeit.default_timer() | |
#turn image into tensor | |
#Transform the input image for use with effnetb2 | |
img=eff_b2_trans(img).unsqueeze(0) | |
#put model into eval mode,make prediction` | |
effnetb2.eval() | |
with torch.inference_mode(): | |
#pass transformed image thorugh the model and turn the prediction logits into probabilities | |
pred_probs=torch.softmax(effnetb2(img),dim=1) | |
#create a prediction label and prediction probability dictionary | |
pred_labels_and_probs={'pizza':pred_probs[0][0], | |
'steak':pred_probs[0][1], | |
'sushi':pred_probs[0][2]} | |
#pred_labels_and_probs=pred_probs | |
#calculate pred time | |
end_time=timeit.default_timer() | |
pred_time=end_time-start_timer | |
#return pred dictionary and pred time | |
return pred_labels_and_probs,f'{pred_time:.4f}' | |
##gradio app | |
#create title,desription and article | |
title="Foodvision Mini" | |
description="An efficientB2 feature extractor" | |
article="created at home" | |
#create example list | |
example_list=['examples/'+i for i in os.listdir('examples')] | |
#create the gradio demo | |
demo=gr.Interface(fn=predict,inputs=gr.Image(type="pil"), | |
outputs=[gr.Label(num_top_classes=3,label="predictions"), | |
gr.Number(label="Prediction time (s)")], | |
examples=example_list, | |
title=title, | |
description=description, | |
article=article) | |
#launch the demo! | |
demo.launch(debug=False,share=True) | |