app.py

import os
from timeit import default_timer as timer
from typing import Tuple, Dict

import gradio as gr
import torch

from model import create_model

title = "Food Vision Mini by Edesak"
desc = "EffitientNetB2 for recognition of Food with 101 classes from [Zero To Mastery Course](https://www.udemy.com/course/pytorch-for-deep-learning/). Used Dataset if Food 101"
article = "My Github page [Edesak](https://github.com/Edesak)"

filename = "labels.txt"

with open(filename, "r") as file:
    class_names = file.read().split("\n")
class_names = class_names[:-1]

model, transform = create_model()
model.load_state_dict(torch.load(f="EffB2_food_big.pth", map_location=torch.device('cpu')))
model.to("cpu")
example_list = [["Examples/" + example] for example in os.listdir("Examples")]


def predict(img) -> Tuple[Dict, float]:
    start_timer = timer()
    img = transform(img).unsqueeze(0)

    model.eval()
    with torch.inference_mode():
        y = torch.softmax(model(img), dim=1)

    pred_labels = {class_names[i]: float(y[0][i]) for i in range(len(class_names))}
    end_time = timer()
    pred_time = round(end_time - start_timer, 4)
    return pred_labels, pred_time


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)")],
                    title=title,
                    description=desc,
                    article=article,
                    examples=example_list)

demo.launch()