import torch
import torchvision
from PIL import Image
import cv2
import numpy as np
import matplotlib.pyplot as plt

import TractionModel as plup

import random
from tqdm import tqdm

import gradio as gr


def snap(image):
    return np.flipud(image)


def init_model(path):
    model = plup.create_model()
    model = plup.load_weights(model, path)
    model.eval()
    return model


def inference(image):
    image = vanilla_transform(image).to(device).unsqueeze(0)
    with torch.no_grad():
        pred = model(image)
    res = float(torch.sigmoid(pred[1].to("cpu")).numpy()[0])
    return {'pull-up': res, 'no pull-up': 1 - res}


norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]
vanilla_transform = torchvision.transforms.Compose([
                                                    torchvision.transforms.Resize(224),
                                                    torchvision.transforms.ToTensor(),
                                                    torchvision.transforms.Normalize(norm_mean, norm_std)])

model = init_model("model-score0.96-f1_10.9-f1_20.99.pt")
if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
model = model.to(device)

iface = gr.Interface(inference, live=True, inputs=gr.inputs.Image(source="upload", tool=None, type='pil'),
                     outputs=gr.outputs.Label())

iface.test_launch()
if __name__ == "__main__":
    iface.launch()