Spaces:
Sleeping
Sleeping
import gradio as gr | |
import os | |
import torch | |
from model import create_effnet_b2_instance | |
from timeit import default_timer as timer | |
from typing import Tuple, Dict | |
# Setup class names | |
with open("class_names.txt", "r") as f: | |
class_names = [food_name.strip() for food_name in f.readlines()] | |
# Create Food101 compatible EffNetB2 instance | |
effnet_transforms,effnetb2_food_101 = create_effnet_b2_instance(num_classes = len(class_names)) | |
# Load the saved model's state_dict() | |
effnetb2_food_101.load_state_dict(torch.load("effnetb2_food101_dict.pth",map_location = torch.device("cpu"))) | |
def predict(img, model = effnetb2_food_101, transforms = effnet_transforms) -> Tuple[Dict,float]: | |
# start a timer | |
start_timer = timer() | |
# transform the image to be used by the model | |
prepreocpressed_image = transforms(img).unsqueeze(0) | |
# turn off regularization and parameters | |
model.eval() | |
with torch.inference_mode(): | |
prediction = model(prepreocpressed_image) | |
probabilities = torch.softmax(prediction,dim = 1) | |
prob_dict = {class_names[i]: float(probabilities[0][i]) for i in range(len(class_names))} | |
# calculate the time | |
end_timer = timer() | |
total_time = end_timer - start_timer | |
return prob_dict,total_time | |
# create the gradio app | |
title = "FoodVision Big Classifier" | |
description = "An EfficientNetB2 feature extractor trained on the Food101 Dataset to classify across 101 possible classes of food." | |
article = "Model created using pytorch" | |
example_list = [["examples/" + example] for example in os.listdir("examples")] | |
# Create Gradio interface | |
demo = gr.Interface( | |
fn=predict, | |
inputs=gr.Image(type="pil"), | |
outputs=[ | |
gr.Label(num_top_classes=5, label="Predictions"), | |
gr.Number(label="Prediction time (s)"), | |
], | |
examples=example_list, | |
title=title, | |
description=description, | |
article=article, | |
) | |
# Launch the app! | |
demo.launch() | |