V1kstrand/FoodClassifier101

1import gradio as gr2import os3import torch4 5from model import create_effnetb2_model6from timeit import default_timer as timer7from typing import Tuple, Dict8 9# get classnames from class_names.txt10with open('class_names.txt', 'r') as f:11    class_names = [food_name.strip() for food_name in f.readlines()]12 13# Create a model14effnetb2, effnetb2_transforms = create_effnetb2_model(num_classes=101)15 16# load saved weights17effnetb2.load_state_dict(18    torch.load(f='09_pretrained_effnetb2_feature_extractor_food101_20perc.pth',19              map_location=torch.device('cpu'),20              ))21 22# Predict function23def predict(img) -> Tuple[Dict, float]:24    """ Trasforms and preforms a prediction on img and returns prediction and time taken25    """26    start_time = timer()27    28    img = effnetb2_transforms(img).unsqueeze(0)29    30    effnetb2.eval()31    with torch.inference_mode():32        pred_probs = torch.softmax(effnetb2(img), dim=1)33        34    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}35    pred_time = round(timer() - start_time, 5)36    37    return pred_labels_and_probs, pred_time38 39 40# Gradio App41title = 'FoodClassifier101'42description = 'An EfficientNetB2 feature extractor computer vision model to classify images of food into 101 different classes'43article = 'Created By David Vikstrand.'44 45example_list = ['sushi.jpg',46                'pizza.jpg']47 48demo = gr.Interface(fn=predict,49                    inputs=gr.Image(type='pil'),50                    outputs=[gr.Label(num_top_classes=5, label='Predictions'),51                             gr.Number(label='Prediction time (s)')],52                    examples=example_list, 53                    title=title,54                    description=description,55                    article=article)56 57demo.launch()58