自存：pytorch版SRGAN中test整个文件夹的图片代码（是由test_image.py改的）

源代码用的是pytorch版本的代码：GitHub - leftthomas/SRGAN: A PyTorch implementation of SRGAN based on CVPR 2017 paper "Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network"

test_dir.py（是由test_image.py改的）

import argparse
import time
import os

import torch
from PIL import Image
from torch.autograd import Variable
from torchvision.transforms import ToTensor, ToPILImage

from model import Generator

parser = argparse.ArgumentParser(description='Test Single Image')
parser.add_argument('--upscale_factor', default=4, type=int, help='super resolution upscale factor')
parser.add_argument('--test_mode', default='GPU', type=str, choices=['GPU', 'CPU'], help='using GPU or CPU')
parser.add_argument('--image_name', type=str, default='crazing_16.png')
parser.add_argument('--model_name', default='netG_epoch_4_199.pth', type=str, help='generator model epoch name')
opt = parser.parse_args()

image_path = './data/train/1x/' # 存放lr图片的文件夹
image_names = os.listdir(image_path)
for image_name in image_names:
    UPSCALE_FACTOR = opt.upscale_factor
    TEST_MODE = True if opt.test_mode == 'GPU' else False
    IMAGE_NAME = image_path + image_name
    MODEL_NAME = opt.model_name

    model = Generator(UPSCALE_FACTOR).eval()
    if TEST_MODE:
        model.cuda()
        model.load_state_dict(torch.load('epochs/' + MODEL_NAME))
    else:
        model.load_state_dict(torch.load('epochs/' + MODEL_NAME, map_location=lambda storage, loc: storage))

    image = Image.open(IMAGE_NAME).convert('RGB')
    image = Variable(ToTensor()(image), volatile=True).unsqueeze(0)
    if TEST_MODE:
        image = image.cuda()

    start = time.clock()
    out = model(image)
    elapsed = (time.clock() - start)
    print('cost' + str(elapsed) + 's')
    out_img = ToPILImage()(out[0].data.cpu())
    # out_img.save('out_srf_' + str(UPSCALE_FACTOR) + '_' + IMAGE_NAME)
    out_img.save('./data/train/4x/' + image_name)