| import gradio as gr | |
| import os | |
| import torch | |
| from model import create_model | |
| from timeit import default_timer as timer | |
| class_names=["Pizza","Steak","Sushi"] | |
| effnetb2,effnetb2_transform=create_model(num_class=3) | |
| effnetb2.load_state_dict(torch.load(f="effnetb2_pizza_steak_sushi (1).pt",map_location=torch.device('cpu'))) | |
| def predict(img): | |
| start_time=timer() | |
| img=effnetb2_transform(img).unsqueeze(0) | |
| effnetb2.eval() | |
| with torch.inference_mode(): | |
| predict_logit=effnetb2(img) | |
| predict_prob=torch.softmax(predict_logit,dim=1) | |
| predict_label=torch.argmax(predict_prob,dim=1) | |
| pred_label_prob={class_names[i]:float(predict_prob[0][i]) for i in range(len(class_names))} | |
| end_time=timer() | |
| pred_time=round(end_time-start_time,4) | |
| return pred_label_prob,pred_time | |
| title="Food Vision Mini 🍕🥩🍣" | |
| description="An EfficientNetB2 feature extractor computer vision model to classify images of food as pizza, steak or sushi." | |
| example_list=[['examples/'+example] for example in os.listdir("examples")] | |
| demo=gr.Interface(fn=predict,inputs=gr.Image(type="pil"),outputs=[gr.Label(num_top_classes=3,label="Prediction"),gr.Number(label="Prediction time")],examples=example_list,title=title,description=description) | |
| demo.launch() | |