import gradio as gr
from PIL import Image
import torch
from torchvision import transforms
from models.models_x import *
import torchvision_x_functional as TF_x
import torchvision.transforms.functional as TF
from torchvision import transforms
import cv2
from timm.models.hub import download_cached_file


cuda = True if torch.cuda.is_available() else False
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
trans = transforms.ToTensor()


LUT0 = Generator3DLUT_identity()
LUT1 = Generator3DLUT_zero()
LUT2 = Generator3DLUT_zero()
classifier = Classifier()
trilinear_ = Tritri()
if cuda:
    LUT0 = LUT0.cuda()
    LUT1 = LUT1.cuda()
    LUT2 = LUT2.cuda()
    classifier = classifier.cuda()

# Load pretrained models
cache = download_cached_file('https://czc-checkpoint.oss-cn-hangzhou.aliyuncs.com/bing/sRGB/LUTs.pth',
                             check_hash=False, progress=True)
LUTs = torch.load(cache, map_location=torch.device('cpu'))
LUT0.load_state_dict(LUTs["0"])
LUT1.load_state_dict(LUTs["1"])
LUT2.load_state_dict(LUTs["2"])
LUT0.eval()
LUT1.eval()
LUT2.eval()

cache = download_cached_file('https://czc-checkpoint.oss-cn-hangzhou.aliyuncs.com/bing/sRGB/classifier.pth',
                             check_hash=False, progress=True)
classifier.load_state_dict(torch.load(cache, map_location=torch.device('cpu')))
classifier.eval()


XLUT0 = Generator3DLUT_identity()
XLUT1 = Generator3DLUT_zero()
XLUT2 = Generator3DLUT_zero()
Xclassifier = Classifier()
Xtrilinear_ = Tritri()
if cuda:
    XLUT0 = XLUT0.cuda()
    XLUT1 = XLUT1.cuda()
    XLUT2 = XLUT2.cuda()
    Xclassifier = Xclassifier.cuda()

# Load pretrained models
cache = download_cached_file('https://czc-checkpoint.oss-cn-hangzhou.aliyuncs.com/bing/XYZ/LUTs.pth',
                             check_hash=False, progress=True)
XLUTs = torch.load(cache, map_location=torch.device('cpu'))
XLUT0.load_state_dict(XLUTs["0"])
XLUT1.load_state_dict(XLUTs["1"])
XLUT2.load_state_dict(XLUTs["2"])
XLUT0.eval()
XLUT1.eval()
XLUT2.eval()

cache = download_cached_file('https://czc-checkpoint.oss-cn-hangzhou.aliyuncs.com/bing/XYZ/classifier.pth',
                             check_hash=False, progress=True)
Xclassifier.load_state_dict(torch.load(cache, map_location=torch.device('cpu')))
Xclassifier.eval()


def generate_LUT(img):
    pred = classifier(img).squeeze()

    LUT = pred[0] * LUT0.LUT + pred[1] * LUT1.LUT + pred[2] * LUT2.LUT  # + pred[3] * LUT3.LUT + pred[4] * LUT4.LUT

    return LUT

def generate_XLUT(img):
    pred = Xclassifier(img).squeeze()

    XLUT = pred[0] * XLUT0.LUT + pred[1] * XLUT1.LUT + pred[2] * XLUT2.LUT  # + pred[3] * LUT3.LUT + pred[4] * LUT4.LUT

    return XLUT


def inference(ori_image, models_n):
    with torch.no_grad():
        if models_n == 'sRGB':
            # img = Image.open(ori_image)
            # img = TF.to_tensor(img).type(Tensor)
            img = trans(ori_image)
            img = img.unsqueeze(0)
            LUT = generate_LUT(img)
            result = trilinear_(LUT, img)
            result = result.permute(0, 3, 1, 2)
            ndarr = result.squeeze().mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
            im = Image.fromarray(ndarr)
        elif models_n == 'XYZ':
            img = trans(ori_image)
            img = img.unsqueeze(0)
            XLUT = generate_XLUT(img)
            result = Xtrilinear_(XLUT, img)
            result = result.permute(0, 3, 1, 2)
            ndarr = result.squeeze().mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
            im = Image.fromarray(ndarr)
    return im


inputs = [gr.inputs.Image(type='pil', label='待增强图片'),
          gr.inputs.Radio(choices=['sRGB', 'XYZ'], type="value", default="sRGB", label="图片色彩空间")]
outputs = [gr.outputs.Image(type='pil', label='增强后图片')]

title = '基于LUT的图像增强演示'

gr.Interface(inference, inputs, outputs, title=title, allow_flagging= 'never',
             examples=[['./examples/example.jpg', 'sRGB']]).launch(enable_queue=True)