def _build_graph(self, inputs):
noisy, gt = inputs
noisy = noisy / 128.0 - 1
gt = gt / 128.0 - 1
def ReluConv2D(name, x, out_channels, use_relu=True):
if use_relu:
x = tf.nn.relu(x, name='%s_relu' % name)
x = Conv2D('%s_conv' % name, x, out_channels)
return x
def ReluDeconv2D(name, x, out_channels):
x = tf.nn.relu(x, name='%s_relu' % name)
x = Deconv2D('%s_deconv' % name, x, out_channels)
return x
with argscope([Conv2D, Deconv2D],
nl=tf.identity,
stride=STRIDE,
kernel_shape=3):
# encoder
e1 = ReluConv2D('enc1', noisy, NF, use_relu=False)
e2 = ReluConv2D('enc2', e1, NF)
e3 = ReluConv2D('enc3', e2, NF)
e4 = ReluConv2D('enc4', e3, NF)
e5 = ReluConv2D('enc5', e4, NF)
# decoder
e6 = ReluDeconv2D('dec1', e5, NF)
e6 = tf.add(e6, e4, name='skip1')
e7 = ReluDeconv2D('dec2', e6, NF)
e8 = ReluDeconv2D('dec3', e7, NF)
e8 = tf.add(e8, e2, name='skip2')
e9 = ReluDeconv2D('dec4', e8, NF)
e10 = ReluDeconv2D('dec5', e9, CHANNELS)
estimate = tf.add(e10, noisy, name="estimate")
self.cost = tf.reduce_mean(tf.squared_difference(estimate, gt),
name="mse")
estimate_scaled = 128.0 * (1.0 + estimate)
gt_scaled = 128.0 * (1.0 + gt)
psnr = symbf.psnr(estimate_scaled, gt_scaled, 255, name="psnr")
print psnr.name
# use tensorboard for visualization
with tf.name_scope("visualization"):
viz = (tf.concat([noisy, estimate, gt], 2) + 1.0) * 128.0
viz = tf.cast(tf.clip_by_value(viz, 0, 255), tf.uint8, name='viz')
tf.summary.image('noisy, estimate, gt',
viz,
max_outputs=max(30, BATCH))
add_moving_summary(self.cost, psnr)
def _build_graph(self, input_vars):
# in [0, 255]
low, high = input_vars
# in [0, 255]
low_ycbcr = tf_rgb2ycbcr(low)
high_ycbcr = tf_rgb2ycbcr(high)
# in [-1, 1]
low_y = tf.expand_dims(low_ycbcr[:, :, :, 0], axis=-1) / 128.0 - 1
high_y = tf.expand_dims(high_ycbcr[:, :, :, 0], axis=-1) / 128.0 - 1
pred_y = (LinearWrap(low_y)
.Conv2D('C1', 64, stride=1, kernel_shape=5, nl=tf.nn.relu)
.Conv2D('C2', 64, stride=1, kernel_shape=3, nl=tf.nn.relu)
.Conv2D('C3', 32, stride=1, kernel_shape=3, nl=tf.nn.relu)
.Conv2D('C4', SCALE ** 2, stride=1, kernel_shape=3, nl=tf.nn.tanh))()
pred_y = pixel_shift("pixshift", pred_y, SCALE, color=False)
self.cost = tf.reduce_mean(tf.squared_difference(pred_y, high_y), name='cost')
# in [0, 255]
scaled_ycbcr = tf.image.resize_bicubic(low_ycbcr, [HR_SIZE_H, HR_SIZE_W])
# [-1, 1] --> [0, 255]
pred_y = 128.0 * (pred_y + 1.)
pred_ycbcr = tf.concat([pred_y, scaled_ycbcr[:, :, :, 1:3]], axis=3)
pred = tf.identity(tf_ycbcr2rgb(pred_ycbcr), name='prediction')
psnr = symbf.psnr(pred, high, 255.)
with tf.name_scope("visualization_rgb"):
low_up = tf.image.resize_bicubic(low, [HR_SIZE_H, HR_SIZE_W]) # bi-linear upscale lowres
viz = (tf.concat([low_up, pred, high], 2))
viz = tf.cast(tf.clip_by_value(viz, 0, 255), tf.uint8, name='viz')
tf.summary.image('input,pred,gt', viz, max_outputs=max(30, BATCH))
low_y_up = (tf.image.resize_bicubic(low_y, [HR_SIZE_H, HR_SIZE_W]) + 1.0) * 128.
high_y_up = (high_y + 1.0) * 128.
with tf.name_scope("visualization_y"):
low_y_up = tf.image.grayscale_to_rgb(low_y_up)
high_y_up = tf.image.grayscale_to_rgb(high_y_up)
pred_y_up = tf.image.grayscale_to_rgb(pred_y)
viz = (tf.concat([low_y_up, pred_y_up, high_y_up], 2))
viz = tf.cast(tf.clip_by_value(viz, 0, 255), tf.uint8, name='viz2')
tf.summary.image('input_y,pred_y,gt_y', viz, max_outputs=max(30, BATCH))
add_moving_summary(self.cost, psnr)
def scaled_psnr(x, y, name):
return symbf.psnr(128. * (x + 1.0), 128. * (y + 1.), 255, name=name)
def scaled_psnr(x, y, name=None):
x = (x + 1.0) * 127.5
y = (y + 1.0) * 127.5
return symbf.psnr(x, y, 255, name=name)