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)
