import os from timeit import default_timer as timer from typing import Tuple from pathlib import Path import gradio as gr import torch from torch import nn from torchvision import transforms from model import create_effnetb2_model class_names = ["pizza", "steak", "sushi"] device = "cpu" # Create model effnetb2, effnetb2_transforms = create_effnetb2_model(num_classes=len(class_names)) # Load saved weights effnetb2.load_state_dict(torch.load("effnetb2.pth", map_location=torch.device(device))) # Define predict function def predict(img: Image) -> Tuple[dict, float]: """Uses EffnetB2 model to transform and predict on img. Returns prediction probabilities and time taken. Args: img (PIL.Image): Image to predict on. Returns: A tuple (pred_labels_and_probs, pred_time), where pred_labels_and_probs is a dict mapping each class name to the probability the model assigns to it, and pred_time is the time taken to predict (in seconds). """ 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, 4) return pred_labels_and_probs, pred_time # Initialize Gradio app title = "FoodVision Mini" description = "EfficientNetB2 feature extractor to classify images of food as pizza, steak, or sushi." article = "From the [Zero to Mastery PyTorch tutorial](https://www.learnpytorch.io/09_pytorch_model_deployment/)" examples = [list(example) for example in Path("examples").glob("*.jpg")] 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, ) demo.lauch()