import torch
import torch.nn as nn
import torchvision
import gradio as gr
from PIL import Image
from torchvision import transforms

agirliklar=torchvision.models.EfficientNet_B2_Weights.DEFAULT
eff_don=agirliklar.transforms()

model=torchvision.models.efficientnet_b2(weights=agirliklar)
model.classifier=nn.Sequential(nn.Linear(1408,100),nn.ReLU(),nn.Linear(100,5))
model.load_state_dict(torch.load('model.pth'))

class_names=['beş', 'bir', 'dört', 'iki', 'üç']

def predict(img):
    """Transforms and performs a prediction on img and returns prediction and time taken.
    """
    # Start the timer
   # img=Image.open(img)
    # Transform the target image and add a batch dimension
    img = eff_don(img).unsqueeze(0)

    # Put model into evaluation mode and turn on inference mode
    model.eval()
    with torch.inference_mode():
        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
        pred_probs = torch.softmax(model(img), dim=1)

    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}


    # Return the prediction dictionary and prediction time
    return pred_labels_and_probs


# Create title, description and article strings
title = "El işaretleri"
description = "Birden beşe kadar olan sayilarin el işaretlerini anlar"

# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
                    inputs=gr.Image(type="pil"), # what are the inputs?
                    outputs=[gr.Label(num_top_classes=5, label="Predictions")], # what are the outputs?
                            # gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
                   # examples=example_list,
                    title=title,
                    description=description)

# Launch the demo!
demo.launch(debug=False, # print errors locally?
            share=True) # generate a publically shareable URL?