def main(argv=None):
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
pascal_voc_lut = pascal_segmentation_lut()
filename_queue = tf.train.string_input_producer([FLAGS.test_data_path],
num_epochs=1)
image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(
filename_queue)
image_batch_tensor = tf.expand_dims(image, axis=0)
annotation_batch_tensor = tf.expand_dims(annotation, axis=0)
input_image_shape = tf.shape(image_batch_tensor)
image_height_width = input_image_shape[1:3]
image_height_width_float = tf.to_float(image_height_width)
image_height_width_multiple = tf.to_int32(
tf.round(image_height_width_float / 32) * 32)
image_batch_tensor = tf.image.resize_images(image_batch_tensor,
image_height_width_multiple)
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
logits = model.model(image_batch_tensor, is_training=False)
pred = tf.argmax(logits, dimension=3)
pred = tf.expand_dims(pred, 3)
pred = tf.image.resize_nearest_neighbor(images=pred,
size=image_height_width)
annotation_batch_tensor = tf.image.resize_nearest_neighbor(
images=annotation_batch_tensor, size=image_height_width)
pred = tf.reshape(pred, [
-1,
])
gt = tf.reshape(annotation_batch_tensor, [
-1,
])
temp = tf.less_equal(gt, FLAGS.num_classes - 1)
weights = tf.cast(temp, tf.int32)
gt = tf.where(temp, gt, tf.cast(temp, tf.uint8))
acc, acc_update_op = tf.contrib.metrics.streaming_accuracy(pred,
gt,
weights=weights)
miou, miou_update_op = tf.contrib.metrics.streaming_mean_iou(
pred, gt, num_classes=FLAGS.num_classes, weights=weights)
with tf.get_default_graph().as_default():
global_vars_init_op = tf.global_variables_initializer()
local_vars_init_op = tf.local_variables_initializer()
init = tf.group(local_vars_init_op, global_vars_init_op)
variable_averages = tf.train.ExponentialMovingAverage(
0.997, global_step)
saver = tf.train.Saver(variable_averages.variables_to_restore())
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
sess.run(init)
ckpt_state = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
model_path = os.path.join(
FLAGS.checkpoint_path,
os.path.basename(ckpt_state.model_checkpoint_path))
print('Restore from {}'.format(model_path))
saver.restore(sess, model_path)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
for i in range(1449):
start = time.time()
image_np, annotation_np, pred_np, tmp_acc, tmp_miou = sess.run(
[image, annotation, pred, acc_update_op, miou_update_op])
_diff_time = time.time() - start
print('{}: cost {:.0f}ms').format(i, _diff_time * 1000)
#upsampled_predictions = pred_np.squeeze()
#plt.imshow(image_np)
#plt.show()
#visualize_segmentation_adaptive(upsampled_predictions, pascal_voc_lut)
acc_res = sess.run(acc)
miou_res = sess.run(miou)
print("Pascal VOC 2012 validation dataset pixel accuracy: " +
str(acc_res))
print("Pascal VOC 2012 validation dataset Mean IoU: " +
str(miou_res))
coord.request_stop()
coord.join(threads)
def main(argv=None):
gpus = range(len(FLAGS.gpu_list.split(',')))
pascal_voc_lut = pascal_segmentation_lut()
class_labels = list(pascal_voc_lut.keys())
with open('data/color_map', 'rb') as f:
color_map = pickle.load(f, encoding="bytes")
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
now = datetime.datetime.now()
StyleTime = now.strftime("%Y-%m-%d-%H-%M-%S")
os.makedirs(FLAGS.logs_path + StyleTime)
if not os.path.exists(FLAGS.checkpoint_path):
os.makedirs(FLAGS.checkpoint_path)
else:
if not FLAGS.restore:
if os.path.exists(FLAGS.checkpoint_path):
shutil.rmtree(FLAGS.checkpoint_path)
os.makedirs(FLAGS.checkpoint_path)
filename_queue = tf.train.string_input_producer([FLAGS.training_data_path],
num_epochs=1000)
image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(
filename_queue)
image, annotation = flip_randomly_left_right_image_with_annotation(
image, annotation)
image = distort_randomly_image_color(image)
image_train_size = [FLAGS.train_size, FLAGS.train_size]
resized_image, resized_annotation = scale_randomly_image_with_annotation_with_fixed_size_output(
image, annotation, image_train_size)
resized_annotation = tf.squeeze(resized_annotation)
image_batch, annotation_batch = tf.train.shuffle_batch(
[resized_image, resized_annotation],
batch_size=FLAGS.batch_size * len(gpus),
capacity=1000,
num_threads=4,
min_after_dequeue=500)
# split
input_images_split = tf.split(image_batch, len(gpus))
input_segs_split = tf.split(annotation_batch, len(gpus))
learning_rate = tf.Variable(FLAGS.learning_rate, trainable=False)
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
# add summary
tf.summary.scalar('learning_rate', learning_rate)
opt = tf.train.AdamOptimizer(learning_rate)
tower_grads = []
reuse_variables = None
for i, gpu_id in enumerate(gpus):
with tf.device('/gpu:%d' % gpu_id):
with tf.name_scope('model_%d' % gpu_id) as scope:
iis = input_images_split[i]
isms = input_segs_split[i]
total_loss, model_loss, output_pred = tower_loss(
iis, isms, class_labels, reuse_variables)
batch_norm_updates_op = tf.group(
*tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope))
reuse_variables = True
grads = opt.compute_gradients(total_loss)
tower_grads.append(grads)
grads = average_gradients(tower_grads)
apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
summary_op = tf.summary.merge_all()
log_image, log_image_data, log_image_name = build_image_summary()
# save moving average
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
# batch norm updates
with tf.control_dependencies(
[variables_averages_op, apply_gradient_op, batch_norm_updates_op]):
train_op = tf.no_op(name='train_op')
saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)
summary_writer = tf.summary.FileWriter(FLAGS.logs_path + StyleTime,
tf.get_default_graph())
# if FLAGS.pretrained_model_path is not None:
# variable_restore_op = slim.assign_from_checkpoint_fn(FLAGS.pretrained_model_path,
# slim.get_trainable_variables(),
# ignore_missing_vars=True)
global_vars_init_op = tf.global_variables_initializer()
local_vars_init_op = tf.local_variables_initializer()
init = tf.group(local_vars_init_op, global_vars_init_op)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
restore_step = 0
if FLAGS.restore:
sess.run(init)
print('continue training from previous checkpoint')
ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
restore_step = int(ckpt.split('.')[0].split('_')[-1])
saver.restore(sess, ckpt)
else:
sess.run(init)
# if FLAGS.pretrained_model_path is not None:
# variable_restore_op(sess)
start = time.time()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
try:
while not coord.should_stop():
for step in range(restore_step, FLAGS.max_steps):
if step != 0 and step % FLAGS.decay_steps == 0:
sess.run(
tf.assign(learning_rate,
learning_rate.eval() * FLAGS.decay_rate))
ml, tl, _ = sess.run([model_loss, total_loss, train_op])
if np.isnan(tl):
print('Loss diverged, stop training')
break
if step % 10 == 0:
avg_time_per_step = (time.time() - start) / 10
start = time.time()
print('Step {:06d}, model loss {:.4f}, total loss {:.4f}, {:.3f} seconds/step, lr: {:.7f}').\
format(step, ml, tl, avg_time_per_step,learning_rate.eval())
if (step + 1) % FLAGS.save_checkpoint_steps == 0:
filename = ('RefineNet' +
'_step_{:d}'.format(step + 1) + '.ckpt')
filename = os.path.join(FLAGS.checkpoint_path,
filename)
saver.save(sess, filename)
print('Write model to: {:s}'.format(filename))
if step % FLAGS.save_summary_steps == 0:
_, tl, summary_str = sess.run(
[train_op, total_loss, summary_op])
summary_writer.add_summary(summary_str,
global_step=step)
if step % FLAGS.save_image_steps == 0:
log_image_name_str = ('%06d' % step)
img_split, seg_split, pred = sess.run(
[iis, isms, output_pred])
img_split = np.squeeze(img_split)[0]
seg_split = np.squeeze(seg_split)[0]
pred = np.squeeze(pred)[0]
#img_split=cv2.resize(img_split,(128,128))
color_seg = np.zeros(
(seg_split.shape[0], seg_split.shape[1], 3))
for i in range(seg_split.shape[0]):
for j in range(seg_split.shape[1]):
color_seg[i, j, :] = color_map[str(
seg_split[i][j])]
color_pred = np.zeros(
(pred.shape[0], pred.shape[1], 3))
for i in range(pred.shape[0]):
for j in range(pred.shape[1]):
color_pred[i,
j, :] = color_map[str(pred[i][j])]
write_img = np.hstack((color_seg, color_pred))
log_image_summary_op = sess.run(log_image,feed_dict={log_image_name: log_image_name_str, \
log_image_data: write_img})
summary_writer.add_summary(log_image_summary_op,
global_step=step)
except tf.errors.OutOfRangeError:
print('finish')
finally:
coord.request_stop()
coord.join(threads)
def main(argv=None):
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
if not os.path.exists(FLAGS.result_path):
os.makedirs(FLAGS.result_path)
filename_queue = tf.train.string_input_producer([FLAGS.test_data_path],
num_epochs=1)
image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(
filename_queue)
image_batch_tensor = tf.expand_dims(image, axis=0)
annotation_batch_tensor = tf.expand_dims(annotation, axis=0)
input_image_shape = tf.shape(image_batch_tensor)
image_height_width = input_image_shape[1:3]
image_height_width_float = tf.to_float(image_height_width)
image_height_width_multiple = tf.to_int32(
tf.round(image_height_width_float / 32) * 32)
image_batch_tensor = tf.image.resize_images(image_batch_tensor,
image_height_width_multiple)
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
logits = model.model(FLAGS.model_type,
image_batch_tensor,
is_training=False)
pred = tf.argmax(logits, dimension=3)
pred = tf.expand_dims(pred, 3)
pred = tf.image.resize_bilinear(images=pred, size=image_height_width)
annotation_batch_tensor = tf.image.resize_bilinear(
images=annotation_batch_tensor, size=image_height_width)
annotation_batch_tensor = tf.div(annotation_batch_tensor, 255)
pred = tf.reshape(pred, [
-1,
])
gt = tf.reshape(annotation_batch_tensor, [
-1,
])
acc, acc_update_op = tf.contrib.metrics.streaming_accuracy(pred, gt)
miou, miou_update_op = tf.contrib.metrics.streaming_mean_iou(
pred, gt, num_classes=FLAGS.num_classes)
with tf.get_default_graph().as_default():
global_vars_init_op = tf.global_variables_initializer()
local_vars_init_op = tf.local_variables_initializer()
init = tf.group(local_vars_init_op, global_vars_init_op)
variable_averages = tf.train.ExponentialMovingAverage(
0.997, global_step)
saver = tf.train.Saver(variable_averages.variables_to_restore())
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)
config = tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False,
gpu_options=gpu_options)
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(init)
ckpt_state = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
model_path = os.path.join(
FLAGS.checkpoint_path,
os.path.basename(ckpt_state.model_checkpoint_path))
print('Restore from {}'.format(model_path))
saver.restore(sess, model_path)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
for i in range(150):
start = time.time()
image_np, annotation_np, pred_np, tmp_acc, tmp_miou = sess.run(
[image, annotation, pred, acc_update_op, miou_update_op])
_diff_time = time.time() - start
print('{}: cost {:.0f}ms').format(i, _diff_time * 1000)
# upsampled_predictions = pred_np.squeeze()
# plt.subplot(131)
# plt.imshow(image_np)
# plt.subplot(132)
# plt.imshow(annotation_np.squeeze(), cmap='gray')
# plt.subplot(133)
# plt.imshow(np.reshape(pred_np, (annotation_np.shape[0], annotation_np.shape[1])).squeeze(), cmap='gray')
# plt.savefig(os.path.join(FLAGS.result_path, str(i) + '.png'))
prediction = np.reshape(
pred_np, (annotation_np.shape[0],
annotation_np.shape[1])).squeeze() * 255
cv2.imwrite(os.path.join(FLAGS.result_path,
str(i) + '.png'), prediction)
print('Test Finished !')
coord.request_stop()
coord.join(threads)
scale_randomly_image_with_annotation_with_fixed_size_output)
slim = tf.contrib.slim
vgg_checkpoint_path = '/home/jochiu/Fully_CNN/vgg_16.ckpt'
log_folder = '/home/jochiu/Fully_CNN/log'
image_train_size = [384, 384]
number_of_classes = 21
tfrecord_filename = 'pascal_augmented_train.tfrecords'
pascal_voc_lut = pascal_segmentation_lut()
class_labels = pascal_voc_lut.keys()
filename_queue = tf.train.string_input_producer([tfrecord_filename],
num_epochs=10)
image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(
filename_queue)
image, annotation = flip_randomly_left_right_image_with_annotation(
image, annotation)
resized_image, resized_annotation = scale_randomly_image_with_annotation_with_fixed_size_output(
image, annotation, image_train_size)
resized_annotation = tf.squeeze(resized_annotation)
image_batch, annotation_batch = tf.train.shuffle_batch(
[resized_image, resized_annotation],
batch_size=1,
capacity=3000,
num_threads=2,
min_after_dequeue=1000)
def main(argv=None):
gpus = range(len(FLAGS.gpu_list.split(',')))
pascal_voc_lut = pascal_segmentation_lut()
class_labels = pascal_voc_lut.keys()
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
now = datetime.datetime.now()
StyleTime = now.strftime("%Y-%m-%d-%H-%M-%S")
os.makedirs(FLAGS.logs_path + StyleTime)
if not os.path.exists(FLAGS.checkpoint_path):
os.makedirs(FLAGS.checkpoint_path)
else:
if not FLAGS.restore:
if os.path.exists(FLAGS.checkpoint_path):
shutil.rmtree(FLAGS.checkpoint_path)
os.makedirs(FLAGS.checkpoint_path)
#input_images = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='input_images')
#input_segs = tf.placeholder(tf.float32, shape=[None, None,None, 1], name='input_segs')
filename_queue = tf.train.string_input_producer([FLAGS.training_data_path],
num_epochs=1000)
image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(
filename_queue)
image, annotation = flip_randomly_left_right_image_with_annotation(
image, annotation)
image = distort_randomly_image_color(image)
image_train_size = [FLAGS.train_size, FLAGS.train_size]
resized_image, resized_annotation = scale_randomly_image_with_annotation_with_fixed_size_output(
image, annotation, image_train_size)
resized_annotation = tf.squeeze(resized_annotation)
image_batch, annotation_batch = tf.train.shuffle_batch(
[resized_image, resized_annotation],
batch_size=FLAGS.batch_size * len(gpus),
capacity=1000,
num_threads=4,
min_after_dequeue=500)
# split
input_images_split = tf.split(image_batch, len(gpus))
input_segs_split = tf.split(annotation_batch, len(gpus))
# 定义损失函数、学习率、滑动平均操作以及训练过程。
learning_rate = tf.Variable(FLAGS.learning_rate, trainable=False)
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
# add summary
tf.summary.scalar('learning_rate', learning_rate)
opt = tf.train.AdamOptimizer(learning_rate)
tower_grads = []
reuse_variables = None
iis = input_images_split[i]
isms = input_segs_split[i]
total_loss, model_loss, output_pred = tower_loss(iis, isms, class_labels,
reuse_variables)
reuse_variables = True
grads = opt.compute_gradients(total_loss)
tower_grads.append(grads)
grads = average_gradients(tower_grads)
apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
summary_op = tf.summary.merge_all()
# save moving average
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
# batch norm updates
with tf.control_dependencies([variables_averages_op, apply_gradient_op]):
train_op = tf.no_op(name='train_op')
saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)
summary_writer = tf.summary.FileWriter(FLAGS.logs_path + StyleTime,
tf.get_default_graph())
if FLAGS.pretrained_model_path is not None:
variable_restore_op = slim.assign_from_checkpoint_fn(
FLAGS.pretrained_model_path,
slim.get_trainable_variables(),
ignore_missing_vars=True)
global_vars_init_op = tf.global_variables_initializer()
local_vars_init_op = tf.local_variables_initializer()
init = tf.group(local_vars_init_op, global_vars_init_op)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
restore_step = 0
if FLAGS.restore:
sess.run(init)
print('continue training from previous checkpoint')
ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
restore_step = int(ckpt.split('.')[0].split('_')[-1])
saver.restore(sess, ckpt)
else:
sess.run(init)
if FLAGS.pretrained_model_path is not None:
variable_restore_op(sess)
start = time.time()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
try:
while not coord.should_stop():
for step in range(restore_step, FLAGS.max_steps):
if step != 0 and step % FLAGS.decay_steps == 0:
sess.run(
tf.assign(learning_rate,
learning_rate.eval() * FLAGS.decay_rate))
ml, tl, _ = sess.run([model_loss, total_loss, train_op])
if np.isnan(tl):
print('Loss diverged, stop training')
break
if step % 10 == 0:
avg_time_per_step = (time.time() - start) / 10
start = time.time()
print('Step {:06d}, model loss {:.6f}, total loss {:.6f}, {:.3f} seconds/step, lr: {:.10f}').\
format(step, ml, tl, avg_time_per_step,learning_rate.eval())
if (step + 1) % FLAGS.save_checkpoint_steps == 0:
filename = ('RefineNet' +
'_step_{:d}'.format(step + 1) + '.ckpt')
filename = os.path.join(FLAGS.checkpoint_path,
filename)
saver.save(sess, filename)
print('Write model to: {:s}'.format(filename))
if step % FLAGS.save_summary_steps == 0:
_, tl, summary_str = sess.run(
[train_op, total_loss, summary_op])
summary_writer.add_summary(summary_str,
global_step=step)
except tf.errors.OutOfRangeError:
print('finish')
finally:
coord.request_stop()
coord.join(threads)
