import os, torch, torchvision, torchvision
import gradio as gr

from model import build_effnetb1
from typing import Dict
from pathlib import Path

# Define Class names
class_names = ["Dark", "Green", "Light", "Medium"]

# Load path for example photo list
exp_list = list(Path("examples/").glob("*.png"))

# Build and load model params
model, transforms = build_effnetb1()
model_path = "effnetb1.pth"
model.load_state_dict(torch.load(f=model_path, map_location="cpu"))


# Predict based on image given | move everything to device("cpu") / Spaces run on CPU
def predict(img) -> Dict:
    # move model to cpu
    model.to("cpu")
    model.eval()
    with torch.inference_mode():
        transformed_image = transforms(img).unsqueeze(dim=0)
        # move input to cpu
        target_image_pred = model(transformed_image.to("cpu"))

    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
    print(target_image_pred_probs)
    pred_labels_and_probs = {class_names[i]: float(target_image_pred_probs[0][i]) for i in range(len(class_names))}

    return pred_labels_and_probs


# Gradio App
title = "Coffee Bean Multi-classifier based on level of roasting ☕️"
description = """Created from multi-classifier model using transfer learning from [EfficientNetB1](https://pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b1.html).
                Model was trained on 10 epochs on default weights, and demonstrated a testing accuarcy of 98%.\n Further information and the source code is provided at my [Github Repo](https://github.com/sehyunlee217/coffee_bean_multi_classification).
                \n\n There are four roasting levels: Green and lightly roasted coffee beans are Laos Typica Bolaven. Doi Chaang are the medium roasted, and Brazil Cerrado are dark roasted. All coffee beans are Arabica beans.\n"""
article = "Dataset from: Ontoum, S., Khemanantakul, T., Sroison, P., Triyason, T., & Watanapa, B. (2022). Coffee Roast Intelligence. arXiv preprint arXiv:2206.01841."

demo = gr.Interface(fn=predict,
                    inputs=gr.Image(type="pil"),
                    outputs=[gr.Label(num_top_classes=4, label="Predictions")],
                    examples=exp_list,
                    title=title,
                    description=description,
                    article=article)

demo.launch()