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)
if pre_threshold:
pred_annotation = utils.argmax_pre_threshold(logits)
pred_annotation = tf.reshape(pred_annotation, [batch_size, IMAGE_SIZE_W, IMAGE_SIZE_L, 1])
else:
pred_annotation = utils.argmax_threshold(logits, threshold)
pred_annotation = tf.reshape(pred_annotation, [batch_size, IMAGE_SIZE_W, IMAGE_SIZE_L, 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, image_test, image_name = sess.run([pred_annotation, image_batch, img_name],
feed_dict={keep_probability: 1.0})
image_name = image_name[0].decode()
image_name = image_name.split('/')[-1].strip('.png')
image_test = np.reshape(image_test, [IMAGE_SIZE_W, IMAGE_SIZE_L, 3])
pred = np.squeeze(pred, axis=3)
pred = np.reshape(pred, [IMAGE_SIZE_W, IMAGE_SIZE_L])
if pre_threshold:
if not os.path.isdir(pre_threshold_dir):
os.makedirs(pre_threshold_dir)
f = h5.File(pre_threshold_dir + image_name + '.h5', 'w')
f['data'] = pred
f.close()
pred = utils.apply_threshold(pred, threshold)
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)
print("Saved no. %s image: %s" % (counter, image_name + '.png'))
coord.request_stop()
coord.join(threads)
