获取DEA清晰图

沈四岁

已于 2023-04-10 17:48:56 修改

阅读量102

点赞数

分类专栏：笔记文章标签： python 开发语言

于 2023-04-10 15:21:16 首次发布

本文链接：https://blog.youkuaiyun.com/swx595182208/article/details/130061147

版权

笔记专栏收录该内容

16 篇文章

订阅专栏

该代码段展示了一个基于PyTorch的图像去雾模型DEA_Net的使用方法。首先，它导入了必要的库，包括Backbone模型和数据转换函数。然后，将图片转换为张量并加载预训练模型到GPU。对输入图像进行补零填充后，通过模型进行处理，最后将输出的清晰图像保存为JPEG文件。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

'''获得清晰图像'''
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")