import os

from torchvision import transforms

import model
import utility
from model.blindsr import BlindSR
import torch
from option import args
import cv2
import numpy as np

from utils import util


def to_img(tensor, arg):
    normalized = tensor[0].data.mul(255 / arg.rgb_range)
    # print(tensor[0])
    # print(normalized)
    out = normalized.byte().permute(1, 2, 0).cpu().numpy()
    # out = np.concatenate((out, out, out), axis=0)
    b = out[:, :, 0].copy()
    g = out[:, :, 1].copy()
    r = out[:, :, 2].copy()

    out[:, :, 0] = r
    out[:, :, 1] = g
    out[:, :, 2] = b
    print(out)

    return out


def main():
    model1 = BlindSR(args)
    model1.load_state_dict(torch.load('experiment/blindsr_x2_bicubic_iso/model/model_600.pt'), strict=False)
    # checkpoint = utility.checkpoint(args)
    #
    # if checkpoint.ok:
    #     model1 = model.Model(args, checkpoint)
    model1 = model1.cuda()
    model1.eval()
    img = cv2.imread('Figs/001 (2).png')
    temp = img[:, :, 0].copy()
    img[:, :, 0] = img[:, :, 2]
    img[:, :, 2] = temp
    transf = transforms.ToTensor()
    img = transf(img)
    img = img.unsqueeze(0).contiguous() * 255
    with torch.no_grad():
        img = img.cuda()
        sr = model1(img)

        sr = utility.quantize(sr, args.rgb_range)
        print(sr)
    result = to_img(sr, args)
    cv2.imwrite('result.jpg', result)


if __name__ == '__main__':
    main()
