from gradio import inputs
import torch
import numpy as np
import torchvision as tv
import gradio as gr

model = tv.models.efficientnet_b0()
num_ftrs = model.classifier[1].in_features
model.classifier[1] = torch.nn.Linear(num_ftrs, 21)
model.load_state_dict(torch.load('model/model_ep=1_acc=0.8909620610367893.pt', map_location = torch.device('cpu')))
model.eval()

classes_to_idx = {'Accenteur mouchet': 0,
 'Bouvreuil pivoine': 1,
 'Chardonneret élégant': 2,
 'Ecureuil roux': 3,
 'Geai des chênes': 4,
 'Grosbec casse-noyaux': 5,
 'Merle noir': 6,
 'Moineau domestique': 7,
 'Moineau friquet': 8,
 'Mésange Nonnette': 9,
 'Mésange bleue': 10,
 'Mésange charbonnière': 11,
 'Mésange huppée': 12,
 'Mésange noire': 13,
 'Pic épeiche': 14,
 'Pinson des arbres': 15,
 'Pinson du Nord': 16,
 'Rougegorge familier': 17,
 'Sittelle torchepot': 18,
 'Tourterelle turque': 19,
 "Verdier d'Europe": 20}

classes = list(classes_to_idx.keys())

preprocess = tv.transforms.Compose([
    tv.transforms.Resize((270, 359)),
    tv.transforms.ToTensor()
    #tv.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

from PIL import Image

def classidy_bird(image):

    inputs = preprocess(image).unsqueeze(0)
    inputs = inputs.to(torch.device('cpu'))

    pred = torch.nn.functional.softmax(model(inputs), dim = 1).detach().numpy()[0]
    return {classes[i] : float(pred[i]) for i in range(21)}

image = gr.inputs.Image(type="pil", shape=(270, 359))
label = gr.outputs.Label(num_top_classes=3)
title = "Poids Plume Classifier"
examples = ['examples/mesange-charbonniere.jpg', 'examples/merle-noir.jpg', 'examples/tourterelle-turque.jpg']

gr.Interface(fn = classidy_bird, inputs=image, outputs=label, capture_session=True, examples=examples, title=title).launch()