AODNet复现: 用gpu批量处理图片

import cv2
import os
import torch
import torch.optim
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt



def dehaze_image(image_addr, name):
    data_hazy = Image.open(image_addr)
    data_hazy = np.array(data_hazy) / 255.0

    data_hazy = torch.from_numpy(data_hazy).float()
    data_hazy = data_hazy.permute(2, 0, 1)
    data_hazy = data_hazy.unsqueeze(0)
    # print(type(data_hazy))
    dehaze_net = torch.load('saved_models/dehaze_net_epoch_17.pth', map_location=torch.device('cuda'))
    data_hazy_cuda = data_hazy.cuda()
    # print(type(data_hazy_cuda))
    # clean_image = dehaze_net(data_hazy).detach().numpy().squeeze()
    clean_image = dehaze_net(data_hazy_cuda).cpu()
    clean_image = clean_image.detach().numpy().squeeze()
    clean_image = np.swapaxes(clean_image, 0, 1)
    clean_image = np.swapaxes(clean_image, 1, 2)

    # save_addr = './results/' + name + '.png'
    save_addr = './results/' + name# ots
    # print('saveaddr:', save_addr)
    # print('yes')
    plt.imsave(save_addr, clean_image)


def dehazeFile(img_Dir):
        img_pathDir = os.listdir(img_Dir)                           # 提取所有文件名,并存在列表中
        print("img_Dir:", img_Dir)                                              # 输出文件路径
        print("img_pathDir:", img_pathDir)                                          # 输出文件名列表
        print(len(img_pathDir))                                     # 输出文件数
        num = 0
        for i in img_pathDir:
            addr = r"./test_images/" + i
            name = i.split("_")[0]
            # dehaze_image(addr, name)
            # print(i)
            dehaze_image(addr, i) #ots
            num = num + 1
            print('handling:{}/{}'.format(num, len(img_pathDir)))
        return


if __name__ == '__main__':
    img_Dir = './test_images'#本地文件路径
    dehazeFile(img_Dir)

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值