import os
import torch
from option import args

import torchvision.transforms as transforms
import cv2 as cv
import numpy as np

# os.environ["CUDA_VISIBLE_DEVICES"] = '8'
#
# img = cv.imread('Figs/512result.bmp')
# print(img.shape)   # numpy数组格式为（H,W,C）

# np_transpose = np.ascontiguousarray(img.transpose((2, 0, 1)))
# img_tensor = torch.from_numpy(np_transpose).float()
# img_tensor.mul_(255 / 255)
# img_tensor = img_tensor.unsqueeze(0).cuda()
# print(img_tensor.size())
from model.blindsr import BlindSR

dasr = BlindSR(args)
dasr.load_state_dict(torch.load('experiment/blindsr_x2_bicubic_iso/model/model_600.pt'), strict=False)
example = torch.ones(1, 3, 96, 96)
dasr.eval()
r1 = dasr(example)
print(r1)
# dasr = dasr.cuda()
# torch.save(dasr, 'experiment/blindsr_x2_bicubic_iso/model/DASR.pt')
# model = torch.load('experiment/blindsr_x2_bicubic_iso/model/DASR.pt')
# with torch.no_grad():
#     sr = dasr(img_tensor)
#     print(sr.data)

# fea = dasr.E(img_tensor, img_tensor)
# sr = dasr.G(img_tensor, fea)

# normalized = sr[0].data.byte()
# ndarr = normalized.permute(1, 2, 0).cpu().numpy()
# # ndarr = np.transpose(ndarr, (1, 2, 0))
# # cv2.cvtColor(ndarr, cv2.COLOR_BGR2)
# b = ndarr[:, :, 0].copy()
# g = ndarr[:, :, 1].copy()
# r = ndarr[:, :, 2].copy()
#
# ndarr[:, :, 0] = r
# ndarr[:, :, 1] = g
# ndarr[:, :, 2] = b
# print(ndarr.shape)
# cv.imwrite('{}{}.png'.format('result', 'SR'), ndarr)

traced_script_module = torch.jit.trace(dasr, example)
r2 = traced_script_module(example)
print(r2)
traced_script_module.save("model.pt")
