'''获得清晰图像'''
from DEA_Net.code.model.backbone import Backbone
def img2tensor(path):
    import torchvision.transforms as transforms
    import cv2 as cv
    img = cv.imread(path)
    transf = transforms.ToTensor()
    img_tensor = transf(img)
    # print('opencv', img)
    # print('torch', img_tensor)
    return img_tensor
def tensor2img(img,name):
    from torchvision import utils as vutils
    vutils.save_image(img, name, normalize=True)
model = Backbone()
model.to('cuda')
ckpt = torch.load('/home8T/swx/yolov3/DEA_Net/trained_models/Hazy4K/PSNR3426_SSIM9885.pth',
                  map_location='cuda')
import torch.nn.functional as F
def pad_img(x, patch_size):
    _, _, h, w = x.size()
    mod_pad_h = (patch_size - h % patch_size) % patch_size
    mod_pad_w = (patch_size - w % patch_size) % patch_size
    x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
    return x
# if isinstance(ckpt, torch.nn.DataParallel):
#     ckpt = ckpt.module
# network.load_state_dict(ckpt.state_dict())

'''用的话直接把下面的解除注释'''
model.load_state_dict(ckpt)
img = img2tensor('/home8T/swx/yolov3/DEA_Net/dataset/Hazy4K/train/hazy/1001_0.89_1.66.png').unsqueeze(0).to('cuda')
# print(img)
H, W = img.shape[2:]
img = pad_img(img, 4)
output = model(img)
tensor2img(output,"./test.jpg")