def get_loader(config, mode='train', pin=False): res_sizes = get_resolutions() _, _, IMAGE_SIZE = get_specified_res(res_sizes, config.phone, config.resolution) # 定义不同模式下的DataLoader if mode == 'train': shuffle = True dataset = LoadData(config.phone, config.dped_dir, IMAGE_SIZE) data_loader = data.DataLoader(dataset=dataset, batch_size=config.batch_size, shuffle=shuffle, num_workers=config.num_thread, pin_memory=pin) else: shuffle = False dataset = LoadData(config.phone, config.dped_dir, IMAGE_SIZE, test=True) data_loader = data.DataLoader(dataset=dataset, batch_size=config.batch_size, shuffle=shuffle, num_workers=config.num_thread, pin_memory=pin) return data_loader
# python test_model.py model=iphone_orig dped_dir=dped/ test_subset=full iteration=all resolution=orig use_gpu=true from scipy import misc import numpy as np import tensorflow as tf from models import resnet import utils import os import sys # process command arguments phone, dped_dir, test_subset, iteration, resolution, use_gpu = utils.process_test_model_args( sys.argv) # get all available image resolutions res_sizes = utils.get_resolutions() # get the specified image resolution IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_SIZE = utils.get_specified_res( res_sizes, phone, resolution) # disable gpu if specified config = tf.ConfigProto( device_count={'GPU': 0}) if use_gpu == "false" else None # create placeholders for input images x_ = tf.placeholder(tf.float32, [None, IMAGE_SIZE]) x_image = tf.reshape(x_, [-1, IMAGE_HEIGHT, IMAGE_WIDTH, 3]) # generate enhanced image enhanced = resnet(x_image)