foodvision_mini / app.py
udaythegreat's picture
Update app.py
cef5e7a verified
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)