#!/usr/bin/env python3
import gradio as gr
import torch
import torchvision
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision import transforms, models
from PIL import Image
import itertools
import numpy as np


args = {
    "model_path": "model_last_epoch_34_torchvision0_3_state.ptw"
}
args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Get classes
dxlabels = ["akiec", "bcc", "bkl", "df", "mel", "nv","vasc"]

# No specific normalization was performed during training
def normtransform(x):
    return x

# Load model
model = torchvision.models.resnet34()
model.fc = torch.nn.Linear(model.fc.in_features, len(dxlabels))
model.load_state_dict(torch.load(args.model_path))
model.eval()
model.to(device=args.device)
torch.set_grad_enabled(False)


def predict(image)->dict:
    global model, dxlabels, args, normtransform
    
    prediction_tensor = torch.zeros([1, len(dxlabels)]).to(device=args.device)

    # Test-time augmentations
    available_sizes = [224]
    target_sizes, hflips, rotations, crops = available_sizes, [0, 1], [0, 90], [0.8]
    aug_combos = [x for x in itertools.product(target_sizes, hflips, rotations, crops)]
    
    # Load image
    img = Image.open(image)
    img = img.convert('RGB')

    # Predict with Test-time augmentation
    for (target_size, hflip, rotation, crop) in tqdm(aug_combos, leave=True):
        tfm = transforms.Compose([
            transforms.Resize(int(target_size // crop)),
            transforms.CenterCrop(target_size),
            transforms.RandomHorizontalFlip(hflip),
            transforms.RandomRotation([rotation, rotation]),
            transforms.ToTensor(),
            # shades_of_grey_torch,
            normtransform
        ])
        test_data = tfm(img).unsqueeze(0).to(device=args.device)
        running_preds = torch.FloatTensor().to(device=args.device)
        outputs = model(test_data)
        prediction_tensor += running_preds
    
    prediction_tensor /= len(aug_combos)
    predictions = F.softmax(prediction_tensor, dim=1)[0].cpu().numpy()
    return {dxlabels[enu]: p for enu, p in enumerate(predictions)}

description = """Research artifact for multi-class predictions of common 
dermatologic tumors. This is the model used in the publication 
[Tschandl P. et al. Nature Medicine 2020](https://www.nature.com/articles/s41591-020-0942-0).

For education and research use only. 
**DO NOT use this to obtain medical advice!**
If you have a skin change in question, seek contact to your physician.
"""

gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="Upload a dermatoscopic image", type="filepath"),
    outputs=gr.outputs.Label(num_top_classes=len(dxlabels)),
    title="Dermatoscopic evaluation",
    description=description,
    allow_flagging="manual"
).launch()