import gradio as gr
from model import create_effnetb2_model
import os
import torch
from torch import nn
from typing import List, Dict, Tuple
from timeit import default_timer as timer

with open('food101_classes.txt', 'r') as f:
    class_names = f.read().splitlines()

model, transforms = create_effnetb2_model(num_classes=len(class_names))

ckpt = torch.load('effnetb2_stepdecay_50epochs.tar', map_location='cpu')
model.load_state_dict(ckpt['model_state_dict'])
model.to('cpu')


def predict(img) -> Tuple[Dict, float]:
    start = timer()
    img = transforms(img)
    img = img.unsqueeze(0)
    img = img.to('cpu')
    model.to('cpu')
    model.eval()
    with torch.inference_mode():
        pred_logits = model(img)
        pred_probs = nn.Softmax(dim=1)(pred_logits).squeeze(0)
        pred_probs_dict = {class_names[i]: pred_probs[i].item() for i in range(len(class_names))}
    end = timer()
    return pred_probs_dict, round(end - start, 4)


examples_dir = 'examples'
examples = [[os.path.join(examples_dir, f)] for f in os.listdir(examples_dir)]

import gradio as gr
title = "Food101 Image Classifier 🥘"
description = "This efficientnetb2 model finetuned on Food101 dataset for 50 epochs with step decay scheduler."
article = "Udemy PyTorch Bootcamp: Created for practice using [Gradio](https://www.gradio.app/)"
demo = gr.Interface(fn=predict,
                    inputs=gr.Image(type="pil", label="Image"),
                    outputs=[gr.Label(label="Predictions", num_top_classes=5),
                             gr.Number(label="Prediction Time (s)")],
                    examples=examples,
                    title=title,
                    description=description,
                    article=article)
demo.launch(share=True)