Spaces:
Sleeping
Sleeping
File size: 5,487 Bytes
24af9b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
### 1. Imports and class names setup ###
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
# Setup class names
class_names = ["pizza", "steak", "sushi"]
### 2. Model and transforms preparation ###
# Create EffNetB2 model
effnetb2, effnetb2_transforms = create_effnetb2_model(
num_classes=3, # len(class_names) would also work
)
# Load saved weights
effnetb2.load_state_dict(
torch.load(
f="09_pretrained_effnetb2_feature_extractor_pizza_steak_sushi_20_percent.pth",
map_location=torch.device("cpu"), # load to CPU
)
)
### 3. Predict function ###
# Create predict function
def predict(img) -> Tuple[Dict, float]:
"""Transforms and performs a prediction on img and returns prediction and time taken.
"""
# Start the timer
start_time = timer()
# Transform the target image and add a batch dimension
img = effnetb2_transforms(img).unsqueeze(0)
# Put model into evaluation mode and turn on inference mode
effnetb2.eval()
with torch.inference_mode():
# Pass the transformed image through the model and turn the prediction logits into prediction probabilities
pred_probs = torch.softmax(effnetb2(img), dim=1)
# Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
# Calculate the prediction time
pred_time = round(timer() - start_time, 5)
# Return the prediction dictionary and prediction time
return pred_labels_and_probs, pred_time
### 4. Gradio app ###
# Create title, description and article strings
title = "FoodVision Mini 🍕🥩🍣"
description = "An EfficientNetB2 feature extractor computer vision model to classify images of food as pizza, steak or sushi."
article = "Created at [09. PyTorch Model Deployment](https://www.learnpytorch.io/09_pytorch_model_deployment/)."
# Create examples list from "examples/" directory
example_list = [["examples/" + example] for example in os.listdir("examples")]
# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
inputs=gr.Image(type="pil"), # what are the inputs?
outputs=[gr.Label(num_top_classes=3, label="Predictions"), # what are the outputs?
gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
# Create examples list from "examples/" directory
examples=example_list,
title=title,
description=description,
article=article)
# Launch the demo!
demo.launch()
# 1. Imports and class name
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
# setup class name
class_names = ['pizza', 'steak', 'sushi']
# 2. Model and transforms preparation
# create effnetb2 model
effnetb2, effnetb2_transforms = create_effnetb2_model(num_classes=3)
# load saved weights
effnetb2.load_state_dict(
torch.load(f='09_pretrained_effnetb2_feature_extractor_pizza_steak_sushi_20_percent',
map_location = torch.device('cpu')
)
)
# 3. predict function
# create predict function
def predict(img) -> Tuple[Dict, float]:
"""
Transforms and performs a prediction on img and returns prediction and time taken.
"""
# start the timer
start_time = timer()
# transform the target img and add a batch dimension
img = effnetb2_transforms(img).unsqueeze(dim=0)
# put model into evaluation mode and turn on inference mode
effnetb2.eval()
with torch.inference_mode():
# pass the transformed image through the model and turn the prediction logits into prediction probabilities
pred_probs = torch.softmax(effnetb2(img), dim=1)
# create a prediction label and prediction probability dictionary for each prediction class
# this is the required format for Gradio's output
pred_labels_and_probs = {class_names[i]: float(pred_probs[0][1]) for i in range(len(class_names))}
# end time
end_time = timer()
# calculate the prediction time
pred_time = round(end_time - start_time, 5)
return pred_labels_and_probs, pred_time
# 4. Gradio App
# Create titme, description, and article strings
title = "FoodVision Mini 🍕🥩🍣"
description = "An EfficientNetB2 feature extractor computer vision model to classify images of food as pizza, steak or sushi."
article = "Learned from [09. PyTorch Model Deployment](https://www.learnpytorch.io/09_pytorch_model_deployment/)."
# create examples list from 'examples/' directory
example_list = [[f'examples/{example}'] for example 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='Prediction'),
gr.Number(label='Prediction time(s)')],
examples=example_list,
title=title,
description=description,
article=article)
# launch the demo!
demo.launch()
|