def main(argv=None):
batch_size = test_batch_size
keep_probability = tf.placeholder(tf.float32, name="keep_probabilty")
image_batch, img_name = tf.train.batch([img, filename_queue[0]],
batch_size)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
logits = FCN4.inference(image_batch, keep_probability)
pred_annotation, pred_annotation_resized = utils.argmax_pre_threshold_resize(
logits)
pred_annotation = tf.reshape(pred_annotation,
[batch_size, IMAGE_HEIGHT, IMAGE_WIDTH, 1])
pred_annotation_resized = tf.reshape(
pred_annotation_resized,
[batch_size, RESIZED_IMAGE_HEIGHT, RESIZED_IMAGE_WIDTH, 1])
print("Setting up Saver...")
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state(saved_dir)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
print("Model restored...")
if not os.path.isdir(test_saved_dir):
os.makedirs(test_saved_dir)
for counter in range(test_batch_num):
pred, pred_resized, image_test, image_name = sess.run(
[pred_annotation, pred_annotation_resized, image_batch, img_name],
feed_dict={keep_probability: 1.0})
image_name = image_name[0].decode()
image_name = image_name.split('/')[-1].strip('.png')
pred = np.squeeze(pred, axis=3)
pred_resized = np.squeeze(pred_resized, axis=3)
image_test = np.reshape(image_test, [IMAGE_HEIGHT, IMAGE_WIDTH, 3])
pred = np.reshape(pred, [IMAGE_HEIGHT, IMAGE_WIDTH])
pred_resized = np.reshape(pred_resized,
[RESIZED_IMAGE_HEIGHT, RESIZED_IMAGE_WIDTH])
if not os.path.isdir(pre_threshold_resize_dir):
os.makedirs(pre_threshold_resize_dir)
f = h5.File(pre_threshold_resize_dir + image_name + '.h5', 'w')
f['data'] = pred
f['data_resized'] = pred_resized
f.close()
pred = utils.apply_threshold(pred, threshold)
pred_resized = utils.apply_threshold(pred_resized, threshold)
pred = sess.run(pred)
pred_resized = sess.run(pred_resized)
utils.save_image(image_test,
test_saved_dir + 'image/',
name=image_name + "_image")
utils.save_image(pred,
test_saved_dir + 'pred/',
name=image_name + "_pred",
category=True)
utils.save_image(pred_resized,
test_saved_dir + 'pred_resized/',
name=image_name + "_pred_resized",
category=True)
print("Saved no. %s image: %s" % (counter, image_name + '.png'))
coord.request_stop()
coord.join(threads)
