### 1. Imports and class names setup ### 
import gradio as gr
import os
import torch
import utils
from typing import Tuple, Dict
from model import TransformerNet
from torchvision import transforms
from PIL import Image
from pathlib import Path

# Get model
model_dir = 'models'
models = list(Path(model_dir).glob("*/*.pth.tar"))
models = sorted(models)

# Get style image
style_dir = 'style_images'
style_list = list(Path(style_dir).glob("*"))
style_list = sorted(style_list)

# Get examples
example_list = [["examples/" + example] for example in os.listdir("examples")]

def transfer(image, model):
    device = 'cpu'

    width = image.size[0]
    height = image.size[1]

    if width > 750 or height > 500:
      iamge = image.thumbnail((712, 474))

    # load model
    style_model = TransformerNet()
    state_dict = torch.load(models[int(model)], map_location=torch.device('cpu'))
    style_model.load_state_dict(state_dict["state_dict"])

    content_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.mul(255))
    ])
    content_image = content_transform(image)
    content_image = content_image.unsqueeze(0).to(device)

    style_model.eval()
    with torch.no_grad():
        style_model.to(device)
        output = style_model(content_image).cpu()

    img = utils.deprocess(output[0])
    img = Image.fromarray(img)
    return img, style_list[int(model)]

# Create title, description and article strings
title = "Image Style Transfer"
description = "Choose a image that you want to transfer and the corresponding style. The app will be transfer your image. You will have received new image."
article = "Model have created base on paper [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://arxiv.org/pdf/1603.08155v1.pdf)."

image_output_1 = gr.Image(label='Tranfer') # output result
image_output_2 = gr.Image(label='Style Image') # Show style image

# Create the Gradio demo
demo = gr.Interface(fn=transfer, # mapping function from input to output
                    inputs=[gr.Image(type="pil", label='Input'),
                            gr.Dropdown(choices=[i.parent.name for i in models], value='rain_princess', type='index', label="Style", info="Chooses kind of style image")], # what are the inputs?
                    outputs=[image_output_1, image_output_2], # our fn has two outputs, therefore we have two outputs
                    label = ['One', "Two"],
                    examples=example_list,
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch()