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)