import os
import torch
import numpy as np

from PIL import Image
import torch.nn as nn
from torch.utils import data

from network import *
from dataset.zurich_night_dataset import zurich_night_DataSet
from configs.test_config import get_arguments


palette = [128, 64, 128, 244, 35, 232, 70, 70, 70, 102, 102, 156, 190, 153, 153, 153, 153, 153, 250, 170, 30,
           220, 220, 0, 107, 142, 35, 152, 251, 152, 70, 130, 180, 220, 20, 60, 255, 0, 0, 0, 0, 142, 0, 0, 70,
           0, 60, 100, 0, 80, 100, 0, 0, 230, 119, 11, 32]
zero_pad = 256 * 3 - len(palette)
for i in range(zero_pad):
    palette.append(0)


def colorize_mask(mask):
    new_mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
    new_mask.putpalette(palette)
    return new_mask


def main():
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    device = torch.device("cuda")


    args = get_arguments()
    if not os.path.exists(args.save):
        os.makedirs(args.save)

    if args.model == 'PSPNet':
        model = PSPNet(num_classes=args.num_classes)
    if args.model == 'DeepLab':
        model = Deeplab(num_classes=args.num_classes)
    if args.model == 'RefineNet':
        model = RefineNet(num_classes=args.num_classes, imagenet=False)

    saved_state_dict = torch.load(args.restore_from)
    model_dict = model.state_dict()
    saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict}
    model_dict.update(saved_state_dict)
    model.load_state_dict(saved_state_dict)

    lightnet = LightNet()
    saved_state_dict = torch.load(args.restore_from_light)
    model_dict = lightnet.state_dict()
    saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict}
    model_dict.update(saved_state_dict)
    lightnet.load_state_dict(saved_state_dict)

    model = model.to(device)
    lightnet = lightnet.to(device)
    model.eval()
    lightnet.eval()

    testloader = data.DataLoader(zurich_night_DataSet(args.data_dir, args.data_list, set=args.set))
    interp = nn.Upsample(size=(1080, 1920), mode='bilinear', align_corners=True)

    weights = torch.log(torch.FloatTensor(
        [0.36869696, 0.06084986, 0.22824049, 0.00655399, 0.00877272, 0.01227341, 0.00207795, 0.0055127, 0.15928651,
         0.01157818, 0.04018982, 0.01218957, 0.00135122, 0.06994545, 0.00267456, 0.00235192, 0.00232904, 0.00098658,
         0.00413907])).cuda()
    weights = (torch.mean(weights) - weights) / torch.std(weights) * args.std + 1.0

    for index, batch in enumerate(testloader):
        if index % 10 == 0:
            print('%d processd' % index)
        image, name = batch
        image = image.to(device)
        with torch.no_grad():
            r = lightnet(image)
            enhancement = image + r
            if args.model == 'RefineNet':
                output2 = model(enhancement)
            else:
                _, output2 = model(enhancement)

        weights_prob = weights.expand(output2.size()[0], output2.size()[3], output2.size()[2], 19)
        weights_prob = weights_prob.transpose(1, 3)
        output2 = output2 * weights_prob
        output = interp(output2).cpu().data[0].numpy()

        output = output.transpose(1,2,0)
        output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)

        output_col = colorize_mask(output)
        output = Image.fromarray(output)

        ###### get the enhanced image
        mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        enhancement = enhancement.cpu().data[0].numpy().transpose(1,2,0)
        enhancement = enhancement * mean_std[1] + mean_std[0]
        enhancement = (enhancement - enhancement.min()) / (enhancement.max()-enhancement.min())
        enhancement = enhancement[:, :, ::-1] * 255  # change to BGR
        enhancement = Image.fromarray(enhancement.astype(np.uint8))

        ###### get the light
        light = r.cpu().data[0].numpy().transpose(1, 2, 0)
        light = (light-light.min()) / (light.max() - light.min())
        light = light[:, :, ::-1] * 255  # change to BGR
        light = Image.fromarray(light.astype(np.uint8))


        name = name[0].split('/')[-1]
        output.save('%s/%s' % (args.save, name))
        output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
        enhancement.save('%s/%s_enhancement.png' % (args.save, name.split('.')[0]))
        light.save('%s/%s_light.png' % (args.save, name.split('.')[0]))


if __name__ == '__main__':
    main()

[((30, 13, 30, 33), 'torch.device', 'torch.device', ({(30, 26, 30, 32): '"""cuda"""'}, {}), "('cuda')", False, 'import torch\n'), ((33, 11, 33, 26), 'configs.test_config.get_arguments', 'get_arguments', ({}, {}), '()', False, 'from configs.test_config import get_arguments\n'), ((44, 23, 44, 52), 'torch.load', 'torch.load', ({(44, 34, 44, 51): 'args.restore_from'}, {}), '(args.restore_from)', False, 'import torch\n'), ((51, 23, 51, 58), 'torch.load', 'torch.load', ({(51, 34, 51, 57): 'args.restore_from_light'}, {}), '(args.restore_from_light)', False, 'import torch\n'), ((63, 13, 63, 80), 'torch.nn.Upsample', 'nn.Upsample', (), '', True, 'import torch.nn as nn\n'), ((34, 11, 34, 36), 'os.path.exists', 'os.path.exists', ({(34, 26, 34, 35): 'args.save'}, {}), '(args.save)', False, 'import os\n'), ((35, 8, 35, 30), 'os.makedirs', 'os.makedirs', ({(35, 20, 35, 29): 'args.save'}, {}), '(args.save)', False, 'import os\n'), ((62, 33, 62, 98), 'dataset.zurich_night_dataset.zurich_night_DataSet', 'zurich_night_DataSet', (), '', False, 'from dataset.zurich_night_dataset import zurich_night_DataSet\n'), ((93, 17, 93, 40), 'PIL.Image.fromarray', 'Image.fromarray', ({(93, 33, 93, 39): 'output'}, {}), '(output)', False, 'from PIL import Image\n'), ((76, 13, 76, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((90, 28, 90, 53), 'numpy.argmax', 'np.argmax', (), '', True, 'import numpy as np\n'), ((65, 24, 68, 21), 'torch.FloatTensor', 'torch.FloatTensor', ({(66, 8, 68, 20): '[0.36869696, 0.06084986, 0.22824049, 0.00655399, 0.00877272, 0.01227341, \n    0.00207795, 0.0055127, 0.15928651, 0.01157818, 0.04018982, 0.01218957, \n    0.00135122, 0.06994545, 0.00267456, 0.00235192, 0.00232904, 0.00098658,\n    0.00413907]'}, {}), '([0.36869696, 0.06084986, 0.22824049, 0.00655399, \n    0.00877272, 0.01227341, 0.00207795, 0.0055127, 0.15928651, 0.01157818, \n    0.04018982, 0.01218957, 0.00135122, 0.06994545, 0.00267456, 0.00235192,\n    0.00232904, 0.00098658, 0.00413907])', False, 'import torch\n'), ((69, 48, 69, 66), 'torch.std', 'torch.std', ({(69, 58, 69, 65): 'weights'}, {}), '(weights)', False, 'import torch\n'), ((69, 15, 69, 34), 'torch.mean', 'torch.mean', ({(69, 26, 69, 33): 'weights'}, {}), '(weights)', False, 'import torch\n')]