| import gradio as gr | |
| import os | |
| import torch | |
| from model import create_effnetb2_model | |
| from timeit import default_timer as timer | |
| from typing import Tuple, Dict | |
| class_names= ['pizza','steak','sushi'] | |
| effnetb2, effnetb2_transforms= create_effnetb2_model(num_classes= 3) | |
| effnetb2.load_state_dict(torch.load(f="09_pretrained_effnetb2_feature_extractor__pizza_steak_sushi_20_percent.pth",map_location= | |
| torch.device("cpu") | |
| )) | |
| def predict(img)-> Tuple[Dict, float]: | |
| start_time= timer() | |
| img= effnetb2_transforms(img).unsqueeze(0) | |
| effnetb2.eval() | |
| with torch.inference_mode(): | |
| pred_probs= torch.softmax(effnetb2(img), dim= 1) | |
| pred_labels_and_probs= {class_names[i]: float(pred_probs[0][i]) for i in range (len(class_names))} | |
| pred_time=round(timer()- start_time, 5) | |
| return pred_labels_and_probs, pred_time | |
| titile= "Foodvision Mini" | |
| description="An efficientnetb2 feature extractor computer vision model to classify images of pizza, steak and sushi." | |
| article= "Created at [09_Pytorch model deployment] (https://www.learnpytorch.io/09_pytorch_model_deployment/)" | |
| 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= "predictions"), | |
| gr.Number(label= "prediction time (s)")], | |
| example_list= example_list, title= title, description= description, article= article) | |
| demo.launch() | |