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
def process():
input_path = generate_random_filename(upload_directory,"jpg")
output_path = generate_random_filename(upload_directory,"jpg")
try:
url = request.json["url"]
# phone: iphone, blackberr or sony
phone = request.json["phone"]
# resolution: orig,high,medium,small,tiny
resolution = request.json["resolution"]
download(url, input_path)
# get the specified image resolution
IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_SIZE = utils.get_specified_res(res_sizes, phone, resolution)
# 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)
with tf.Session(config=config) as sess:
saver = tf.train.Saver()
saver.restore(sess, "models_orig/" + phone + "_orig")
image = np.float16(misc.imresize(misc.imread(filename), res_sizes[phone])) / 255
image_crop = utils.extract_crop(image, resolution, phone, res_sizes)
image_crop_2d = np.reshape(image_crop, [1, IMAGE_SIZE])
enhanced_2d = sess.run(enhanced, feed_dict={x_: image_crop_2d})
enhanced_image = np.reshape(enhanced_2d, [IMAGE_HEIGHT, IMAGE_WIDTH, 3])
misc.imsave(filename, enhanced_image)
callback = send_file(output_path, mimetype='image/jpeg')
return callback, 200
except:
traceback.print_exc()
return {'message': 'input error'}, 400
finally:
clean_all([
input_path,
output_path
])
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)
with tf.Session(config=config) as sess:
test_dir = dped_dir + 'Test/' + phone.replace("_orig", "") + '/Full/'