image_u = tf.placeholder(dtype=tf.float32,
shape=[None, 256, 256, 3],
name='image_u')
# correct image
image_r = tf.placeholder(dtype=tf.float32,
shape=[None, 256, 256, 3],
name='image_r')
training_flag = tf.placeholder(tf.bool)
learning_rate = tf.placeholder(tf.float32, shape=[], name='learning_rate')
lr_sum = tf.summary.scalar('lr', learning_rate)
# generated color image by u-net
U_NET = UNet(input_=image_u, real_=image_r, is_training=training_flag)
gen_image = U_NET.u_net(inputs=image_u, training=training_flag)
G_sum = tf.summary.image("gen_image", gen_image, max_outputs=10)
# loss of u-net
errG = U_NET.l1_loss(gt=image_r, gen=gen_image)
# errG = U_NET.mse_loss(gt=image_r, gen=gen_image)
# errG = U_NET.ssim_loss(gt=image_r, gen=gen_image)
# errG = U_NET.msssim_loss(gt=image_r, gen=gen_image)
# errG = U_NET.gdl_loss(gt=image_r, gen=gen_image)
# errG = U_NET.l2_l1_loss(gt=image_r, gen=gen_image, alpha=0.8)
# errG = U_NET.ssim_l1_loss(gt=image_r, gen=gen_image, alpha=0.8)
# errG = U_NET.msssim_l1_loss(gt=image_r, gen=gen_image, alpha=0.8)
# errG = U_NET.gdl_l1_loss(gt=image_r, gen=gen_image, alpha=0.8)
errG_sum = tf.summary.scalar("loss", errG)
t_var = tf.trainable_variables()
if __name__ == '__main__':
# underwater image
image_u = tf.placeholder(dtype=tf.float32, shape=[1, 256, 256, 3], name='image_u')
# correct image
image_r = tf.placeholder(dtype=tf.float32, shape=[1, 256, 256, 3], name='image_r')
# load test image
test_image = normalize_image(misc.imresize(misc.imread(image_path), size=(256, 256), interp='cubic'))
print(test_image)
real_image = normalize_image(misc.imresize(misc.imread(gt_image_path), size=(256, 256), interp='cubic'))
print(real_image)
test_image_np = np.empty(shape=[1, 256, 256, 3], dtype=np.float32)
test_image_np[0, :, :, :] = test_image
test_image_tf = tf.convert_to_tensor(test_image_np)
# laod model
U_NET = UNet(input_=image_u, real_=image_r, is_training=False)
gen_image = U_NET.u_net(inputs=test_image_tf, training=False)
# load weight
saver = tf.train.Saver(max_to_keep=1)
ckpt = tf.train.get_checkpoint_state(ckpt_path)
# generating
begin_time = time.time()
with tf.Session() as sess:
saver.restore(sess=sess, save_path=ckpt.model_checkpoint_path)
# gen_image = U_NET.u_net(inputs=test_image, training=False)
gen = np.asarray(sess.run(gen_image), dtype=np.float32)
print(gen.shape)
print(gen[0])
end_time = time.time()
print("Time cost: %f" % (end_time - begin_time))
misc.imsave('./res_gen2.png', gen[0])
os.makedirs(gen_path)
print("Params Config:\n")
print(" Batch Size: %f" % test_batch_size)
# underwater image
image_u = tf.placeholder(dtype=tf.float32,
shape=[None, 256, 256, 3],
name='image_u')
# correct image
image_r = tf.placeholder(dtype=tf.float32,
shape=[None, 256, 256, 3],
name='image_r')
# load model
U_NET = UNet(input_=image_u, real_=image_r, is_training=False)
gen_images = U_NET.u_net(inputs=image_u, training=False)
# load image
test_path = np.asarray(glob.glob(test_path + '*.png'))
test_path.sort()
# process image
num_test_image = len(test_path)
test_x = np.empty(shape=[num_test_image, 256, 256, 3], dtype=np.float32)
i = 0
for path_i in test_path:
img_i = normalize_image(
misc.imresize(misc.imread(path_i).astype('float32'),
size=(256, 256),
interp='cubic'))
# img_i = normalize_image(misc.imread(path_i).astype('float32'))
test_x[i, :, :, :] = img_i
