def evaluate_model(path, model, pixel_mean, pixel_std, SR_FACTOR=3, sigma=1):
"""
Computes average Peak Signal to Noise Ratio over a set of images and
the versions super resolved by a srcnn or variant.
Args:
path (string): relative path to directory containing images for evaluation
model (PyTorch model): the model to be evaluated
pixel_mean (float): mean luminance value to be used for standardization
pixel_std (float): std. dev. of luminance value to be used for standardization
SR_FACTOR (int): super resolution factor
sigma (int): the std. dev. to use for the gaussian blur
"""
img_names = [im for im in os.listdir(path) if im[-4:] == '.bmp' or im[-4:] == '.jpg']
blurred_img_psnrs = []
out_img_psnrs = []
for test_im in img_names:
blurred_test_im = distort_image(path=path+test_im, factor=SR_FACTOR, sigma=sigma)
ImageFile = Image.open(path+test_im)
im = np.array(ImageFile.convert('YCbCr'))
#normalize
model_input = blurred_test_im[:, :, 0] / 255.0
#standardize
model_input -= pixel_mean
model_input /= pixel_std
im_out_Y = model(torch.tensor(model_input,
dtype=torch.float).unsqueeze(0).unsqueeze(0).to(DEVICE))
im_out_Y = im_out_Y.detach().squeeze().squeeze().cpu().numpy().astype(np.float64)
im_out_viz = np.zeros((im_out_Y.shape[0], im_out_Y.shape[1], 3))
#unstandardize
im_out_Y = (im_out_Y * pixel_std) + pixel_mean
#un-normalize
im_out_Y *= 255.0
im_out_viz[:, :, 0] = im_out_Y
im_out_viz[:, :, 1] = im[:, :, 1]
im_out_viz[:, :, 2] = im[:, :, 2]
im_out_viz[:, :, 0] = np.around(im_out_viz[:, :, 0])
#psnr on Y only
blur_psnr = psnr(im[:, :, 0], blurred_test_im[:, :, 0])
sr_psnr = psnr(im[:, :, 0], im_out_viz[:, :, 0])
blurred_img_psnrs.append(blur_psnr)
out_img_psnrs.append(sr_psnr)
mean_blur_psnr = np.mean(np.array(blurred_img_psnrs))
mean_sr_psnr = np.mean(np.array(out_img_psnrs))
return mean_blur_psnr, mean_sr_psnr
def build_valid_rl(self, shuffle=False):
print("-" * 80)
print("Build valid graph on shuffled data")
with tf.device("/cpu:0"):
# shuffled valid data: for choosing validation model
self.iterators['valid_shuffle'] = self.datasets['valid'].make_one_shot_iterator()
x_valid_shuffle, y_valid_shuffle = self.iterators['valid_shuffle'].get_next()
x_bicubic = tf.image.resize_bicubic(
x_valid_shuffle, [128, 128])
output = self._model(x_valid_shuffle, is_training=True, reuse=True)
self.valid_shuffle_psnr = psnr(y_valid_shuffle, output) - psnr(y_valid_shuffle, x_bicubic)
def _build_valid(self):
if self.x_valid is not None:
print("-" * 80)
print("Build valid graph")
output = self._model(self.x_valid, False, reuse=True)
self.valid_psnr = psnr(self.y_valid, output)
def _build_train(self):
print("-" * 80)
print("Build train graph")
output = self._model(self.x_train, is_training=True)
target = (self.y_train - 127) / 127
self.loss = tf.reduce_mean(
tf.losses.absolute_difference(target, output))
train_loss = self.loss
self.train_psnr = psnr(self.y_train, output)
tf.summary.scalar('loss', self.loss)
output = output * 127 + 127
output = tf.clip_by_value(output, 0, 255)
input_img = self.x_train*127 + 127
bicubic_img = tf.image.resize_bicubic(input_img, [128, 128])
tf.summary.image("output", tf.cast(output, tf.uint8))
tf.summary.image("target", tf.cast(self.y_train, tf.uint8))
tf.summary.image("input", tf.cast(input_img, tf.uint8))
tf.summary.image("bicubic", tf.cast(bicubic_img, tf.uint8))
tf_variables = [
var for var in tf.trainable_variables() if (
var.name.startswith(self.name) and "aux_head" not in var.name)]
self.num_vars = count_model_params(tf_variables)
print("Model has {0} params".format(self.num_vars))
self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
train_loss,
tf_variables,
self.global_step,
clip_mode=self.clip_mode,
grad_bound=self.grad_bound,
l2_reg=self.l2_reg,
lr_init=self.lr_init,
lr_dec_start=self.lr_dec_start,
lr_dec_every=self.lr_dec_every,
lr_dec_rate=self.lr_dec_rate,
lr_cosine=self.lr_cosine,
lr_max=self.lr_max,
lr_min=self.lr_min,
lr_T_0=self.lr_T_0,
lr_T_mul=self.lr_T_mul,
num_train_batches=self.num_train_batches,
optim_algo=self.optim_algo,
sync_replicas=self.sync_replicas,
num_aggregate=self.num_aggregate,
num_replicas=self.num_replicas)
tf.summary.scalar('lr', self.lr)
self.summaries = tf.summary.merge_all()
def build_valid_rl(self, shuffle=False):
print("-" * 80)
print("Build valid graph on shuffled data")
with tf.device("/cpu:0"):
# shuffled valid data: for choosing validation model
if not shuffle and self.data_format == "NCHW":
self.inputs["valid_original"] = np.transpose(
self.inputs["valid_original"], [0, 3, 1, 2])
self.iterators['valid_shuffle'] = self.datasets[
'valid'].make_one_shot_iterator()
x_valid_shuffle, y_valid_shuffle = self.iterators[
'valid_shuffle'].get_next()
output = self._model(x_valid_shuffle, is_training=True, reuse=True)
self.valid_shuffle_psnr = psnr(y_valid_shuffle, output)
def _build_test(self):
print("-" * 80)
print("Build test graph")
output = self._model(self.x_test, False, reuse=True)
self.test_psnr = psnr(self.y_test, output)
