def train(resume_step=None):
# < preparing arguments >
if FLAGS.float_type == 16:
print('\n< using tf.float16 >\n')
float_type = tf.float16
else:
print('\n< using tf.float32 >\n')
float_type = tf.float32
new_layer_names = FLAGS.new_layer_names
if FLAGS.new_layer_names is not None:
new_layer_names = new_layer_names.split(',')
# < data set >
data_list = FLAGS.subsets_for_training.split(',')
if len(data_list) < 1:
data_list = ['train']
list_images = []
list_labels = []
with tf.device('/cpu:0'):
reader = SegmentationImageReader(
FLAGS.database,
data_list, (FLAGS.train_image_size, FLAGS.train_image_size),
FLAGS.random_scale,
random_mirror=True,
random_blur=True,
random_rotate=FLAGS.random_rotate,
color_switch=FLAGS.color_switch,
scale_rate=(FLAGS.scale_min, FLAGS.scale_max))
for _ in xrange(FLAGS.gpu_num):
image_batch, label_batch = reader.dequeue(FLAGS.batch_size)
list_images.append(image_batch)
list_labels.append(label_batch)
# < network >
model = pspnet_mg.PSPNetMG(
reader.num_classes,
mode='train',
resnet=FLAGS.network,
bn_mode='frozen' if FLAGS.bn_frozen else 'gather',
data_format=FLAGS.data_format,
initializer=FLAGS.initializer,
fine_tune_filename=FLAGS.fine_tune_filename,
wd_mode=FLAGS.weight_decay_mode,
gpu_num=FLAGS.gpu_num,
float_type=float_type,
has_aux_loss=FLAGS.has_aux_loss,
train_like_in_paper=FLAGS.train_like_in_paper,
structure_in_paper=FLAGS.structure_in_paper,
new_layer_names=new_layer_names,
loss_type=FLAGS.loss_type,
consider_dilated=FLAGS.consider_dilated)
train_ops = model.build_train_ops(list_images, list_labels)
# < log dir and model id >
logdir = LogDir(FLAGS.database, model_id())
logdir.print_all_info()
if not os.path.exists(logdir.log_dir):
print('creating ', logdir.log_dir, '...')
os.mkdir(logdir.log_dir)
if not os.path.exists(logdir.database_dir):
print('creating ', logdir.database_dir, '...')
os.mkdir(logdir.database_dir)
if not os.path.exists(logdir.exp_dir):
print('creating ', logdir.exp_dir, '...')
os.mkdir(logdir.exp_dir)
if not os.path.exists(logdir.snapshot_dir):
print('creating ', logdir.snapshot_dir, '...')
os.mkdir(logdir.snapshot_dir)
gpu_options = tf.GPUOptions(allow_growth=False)
config = tf.ConfigProto(log_device_placement=False,
gpu_options=gpu_options,
allow_soft_placement=True)
sess = tf.Session(config=config)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
init = [
tf.global_variables_initializer(),
tf.local_variables_initializer()
]
sess.run(init)
# < convert npy to .ckpt >
step = 0
if '.npy' in FLAGS.fine_tune_filename:
# This can transform .npy weights with variables names being the same to the tf ckpt model.
fine_tune_variables = []
npy_dict = np.load(FLAGS.fine_tune_filename).item()
new_layers_names = ['Momentum']
for v in tf.global_variables():
if any(elem in v.name for elem in new_layers_names):
continue
name = v.name.split(':0')[0]
if name not in npy_dict:
continue
v.load(npy_dict[name], sess)
fine_tune_variables.append(v)
saver = tf.train.Saver(var_list=fine_tune_variables)
saver.save(sess, logdir.snapshot_dir + '/model.ckpt', global_step=0)
return
# < load pre-trained model>
import_variables = tf.trainable_variables()
if FLAGS.fine_tune_filename is not None and resume_step is None:
fine_tune_variables = []
new_layers_names = model.new_layers_names
new_layers_names.append('Momentum')
new_layers_names.append('up_sample')
for v in import_variables:
if any(elem in v.name for elem in new_layers_names):
print('< Finetuning Process: not import %s >' % v.name)
continue
fine_tune_variables.append(v)
loader = tf.train.Saver(var_list=fine_tune_variables, allow_empty=True)
loader.restore(sess, FLAGS.fine_tune_filename)
print('< Succesfully loaded fine-tune model from %s. >' %
FLAGS.fine_tune_filename)
elif resume_step is not None:
# ./snapshot/model.ckpt-3000
i_ckpt = logdir.snapshot_dir + '/model.ckpt-%d' % resume_step
loader = tf.train.Saver(max_to_keep=0)
loader.restore(sess, i_ckpt)
step = resume_step
print('< Succesfully loaded model from %s at step=%s. >' %
(i_ckpt, resume_step))
else:
print('< Not import any model. >')
f_log = open(logdir.exp_dir + '/' + str(datetime.datetime.now()) + '.txt',
'w')
f_log.write('step,loss,precision,wd\n')
f_log.write(sorted_str_dict(FLAGS.__dict__) + '\n')
print('\n< training process begins >\n')
average_loss = 0.0
show_period = 20
snapshot = FLAGS.snapshot
max_iter = FLAGS.train_max_iter
lrn_rate = FLAGS.lrn_rate
lr_step = []
if FLAGS.lr_step is not None:
temps = FLAGS.lr_step.split(',')
for t in temps:
lr_step.append(int(t))
saver = tf.train.Saver(max_to_keep=2)
t0 = None
wd_rate = FLAGS.weight_decay_rate
wd_rate2 = FLAGS.weight_decay_rate2
if FLAGS.save_first_iteration == 1:
saver.save(sess, logdir.snapshot_dir + '/model.ckpt', global_step=step)
has_nan = False
while step < max_iter + 1:
if FLAGS.poly_lr == 1:
lrn_rate = ((1 - 1.0 * step / max_iter)**0.9) * FLAGS.lrn_rate
step += 1
if len(lr_step) > 0 and step == lr_step[0]:
lrn_rate *= FLAGS.step_size
lr_step.remove(step)
_, loss, wd, precision = sess.run(
[train_ops, model.loss, model.wd, model.precision_op],
feed_dict={
model.lrn_rate_ph: lrn_rate,
model.wd_rate_ph: wd_rate,
model.wd_rate2_ph: wd_rate2
})
if math.isnan(loss) or math.isnan(wd):
print('\nloss or weight norm is nan. Training Stopped!\n')
has_nan = True
break
average_loss += loss
if step % snapshot == 0:
saver.save(sess,
logdir.snapshot_dir + '/model.ckpt',
global_step=step)
sess.run([tf.local_variables_initializer()])
if step % show_period == 0:
left_hours = 0
if t0 is not None:
delta_t = (datetime.datetime.now() - t0).total_seconds()
left_time = (max_iter - step) / show_period * delta_t
left_hours = left_time / 3600.0
t0 = datetime.datetime.now()
average_loss /= show_period
f_log.write('%d,%f,%f,%f\n' % (step, average_loss, precision, wd))
f_log.flush()
print('%s %s] Step %s, lr = %f, wd_rate = %f, wd_rate_2 = %f ' \
% (str(datetime.datetime.now()), str(os.getpid()), step, lrn_rate, wd_rate, wd_rate2))
print('\t loss = %.4f, precision = %.4f, wd = %.4f' %
(average_loss, precision, wd))
print('\t estimated time left: %.1f hours. %d/%d' %
(left_hours, step, max_iter))
average_loss = 0.0
coord.request_stop()
coord.join(threads)
return f_log, logdir, has_nan # f_log and logdir returned for eval.
def train(resume_step=None):
global_step = tf.get_variable('global_step', [], dtype=tf.int64,
initializer=tf.constant_initializer(0), trainable=False)
image_size = FLAGS.train_image_size
print '================',
if FLAGS.data_type == 16:
print 'using tf.float16 ====================='
data_type = tf.float16
print 'can not use float16 at this moment, because of tf.nn.bn, if using fused_bn, the learning will be nan',
print ', no idea what happened.'
else:
print 'using tf.float32 ====================='
data_type = tf.float32
data_list = FLAGS.subsets_for_training.split(',')
if len(data_list) < 1:
data_list = ['train']
print data_list
images = []
labels = []
with tf.device('/cpu:0'):
reader = SegmentationImageReader(
FLAGS.server,
FLAGS.database,
data_list,
(image_size, image_size),
FLAGS.random_scale,
random_mirror=True,
random_blur=True,
random_rotate=FLAGS.random_rotate,
color_switch=FLAGS.color_switch,
scale_rate=(FLAGS.scale_min, FLAGS.scale_max))
print '================ Database Info ================'
for i in range(FLAGS.gpu_num):
with tf.device('/cpu:0'):
image_batch, label_batch = reader.dequeue(FLAGS.batch_size)
images.append(image_batch)
labels.append(label_batch)
wd_rate_ph = tf.placeholder(data_type, shape=())
wd_rate2_ph = tf.placeholder(data_type, shape=())
lrn_rate_ph = tf.placeholder(data_type, shape=())
new_layer_names = FLAGS.new_layer_names
if FLAGS.new_layer_names is not None:
new_layer_names = new_layer_names.split(',')
assert 'pspnet' in FLAGS.network
resnet = 'resnet_v1_101'
PSPModel = pspnet_mg.PSPNetMG
with tf.variable_scope(resnet):
model = PSPModel(reader.num_classes, lrn_rate_ph, wd_rate_ph, wd_rate2_ph,
mode='train', bn_epsilon=FLAGS.epsilon, resnet=resnet,
norm_only=FLAGS.norm_only,
initializer=FLAGS.initializer,
fix_blocks=FLAGS.fix_blocks,
fine_tune_filename=FLAGS.fine_tune_filename,
bn_ema=FLAGS.ema_decay,
bn_frozen=FLAGS.bn_frozen,
wd_mode=FLAGS.weight_decay_mode,
fisher_filename=FLAGS.fisher_filename,
gpu_num=FLAGS.gpu_num,
float_type=data_type,
fisher_epsilon=FLAGS.fisher_epsilon,
has_aux_loss=FLAGS.has_aux_loss,
train_like_in_paper=FLAGS.train_like_in_paper,
structure_in_paper=FLAGS.structure_in_paper,
new_layer_names=new_layer_names,
loss_type=FLAGS.loss_type,
train_conv2dt=FLAGS.train_conv2dt)
model.inference(images)
model.build_train_op(labels)
names = []
num_params = 0
for v in tf.trainable_variables():
# print v.name
names.append(v.name)
num = 1
for i in v.get_shape().as_list():
num *= i
num_params += num
print "Trainable parameters' num: %d" % num_params
print 'iou precision shape: ', model.predictions.get_shape(), labels[0].get_shape()
pred = tf.reshape(model.predictions, [-1, ])
gt = tf.reshape(labels[0], [-1, ])
indices = tf.squeeze(tf.where(tf.less_equal(gt, reader.num_classes - 1)), 1)
gt = tf.cast(tf.gather(gt, indices), tf.int32)
pred = tf.gather(pred, indices)
precision_op, update_op = tf.contrib.metrics.streaming_mean_iou(pred, gt, num_classes=reader.num_classes)
# ========================= end of building model ================================
step = 0
logdir = LogDir(FLAGS.database, FLAGS.log_dir, FLAGS.weight_decay_mode)
logdir.print_all_info()
if not os.path.exists(logdir.log_dir):
print 'creating ', logdir.log_dir, '...'
os.mkdir(logdir.log_dir)
if not os.path.exists(logdir.database_dir):
print 'creating ', logdir.database_dir, '...'
os.mkdir(logdir.database_dir)
if not os.path.exists(logdir.exp_dir):
print 'creating ', logdir.exp_dir, '...'
os.mkdir(logdir.exp_dir)
if not os.path.exists(logdir.snapshot_dir):
print 'creating ', logdir.snapshot_dir, '...'
os.mkdir(logdir.snapshot_dir)
init = [tf.global_variables_initializer(), tf.local_variables_initializer()]
gpu_options = tf.GPUOptions(allow_growth=False)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options, allow_soft_placement=True)
sess = tf.Session(config=config)
sess.run(init)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
if '.npy' in FLAGS.fine_tune_filename:
# This can transform .npy weights with variables names being the same to the tf ckpt model.
fine_tune_variables = []
npy_dict = np.load(FLAGS.fine_tune_filename).item()
new_layers_names = ['Momentum']
for v in tf.global_variables():
print '=====Saving initial snapshot process:',
if any(elem in v.name for elem in new_layers_names):
print 'not import', v.name
continue
name = v.name.split(':0')[0]
if name not in npy_dict:
print 'not find', v.name
continue
v.load(npy_dict[name], sess)
print 'saving', v.name
fine_tune_variables.append(v)
saver = tf.train.Saver(var_list=fine_tune_variables)
saver.save(sess, logdir.snapshot_dir + '/model.ckpt', global_step=0)
return
import_variables = tf.trainable_variables()
if FLAGS.fix_blocks > 0 or FLAGS.bn_frozen > 0:
import_variables = tf.global_variables()
if FLAGS.fine_tune_filename is not None and resume_step is None:
fine_tune_variables = []
new_layers_names = model.new_layers_names
new_layers_names.append('Momentum')
new_layers_names.append('up_sample')
for v in import_variables:
if any(elem in v.name for elem in new_layers_names):
print '=====Finetuning Process: not import %s' % v.name
continue
fine_tune_variables.append(v)
loader = tf.train.Saver(var_list=fine_tune_variables, allow_empty=True)
loader.restore(sess, FLAGS.fine_tune_filename)
print('=====Succesfully loaded fine-tune model from %s.' % FLAGS.fine_tune_filename)
elif resume_step is not None:
# ./snapshot/model.ckpt-3000
i_ckpt = logdir.snapshot_dir + '/model.ckpt-%d' % resume_step
loader = tf.train.Saver(max_to_keep=0)
loader.restore(sess, i_ckpt)
step = resume_step
print('=====Succesfully loaded model from %s at step=%s.' % (i_ckpt, resume_step))
else:
print '=====Not import any model.'
print '=========================== training process begins ================================='
f_log = open(logdir.exp_dir + '/' + str(datetime.datetime.now()) + '.txt', 'w')
f_log.write('step,loss,precision,wd\n')
f_log.write(sorted_str_dict(FLAGS.__dict__) + '\n')
average_loss = 0.0
show_period = 20
snapshot = FLAGS.snapshot
max_iter = FLAGS.train_max_iter
lrn_rate = FLAGS.lrn_rate
lr_step = []
if FLAGS.lr_step is not None:
temps = FLAGS.lr_step.split(',')
for t in temps:
lr_step.append(int(t))
# fine_tune_variables = []
# for v in tf.global_variables():
# if 'Momentum' in v.name:
# continue
# print '=====Saving initial snapshot process: saving %s' % v.name
# fine_tune_variables.append(v)
#
# saver = tf.train.Saver(var_list=fine_tune_variables)
# saver.save(sess, logdir.snapshot_dir + '/model.ckpt', global_step=0)
saver = tf.train.Saver(max_to_keep=2)
t0 = None
wd_rate = FLAGS.weight_decay_rate
wd_rate2 = FLAGS.weight_decay_rate2
if FLAGS.save_first_iteration == 1:
saver.save(sess, logdir.snapshot_dir + '/model.ckpt', global_step=step)
has_nan = False
while step < max_iter + 1:
if FLAGS.poly_lr == 1:
lrn_rate = ((1-1.0*step/max_iter)**0.9) * FLAGS.lrn_rate
step += 1
if len(lr_step) > 0 and step == lr_step[0]:
lrn_rate *= FLAGS.step_size
lr_step.remove(step)
_, loss, wd, update, precision = sess.run([
model.train_op, model.loss, model.wd, update_op, precision_op
],
feed_dict={
lrn_rate_ph: lrn_rate,
wd_rate_ph: wd_rate,
wd_rate2_ph: wd_rate2
}
)
if math.isnan(loss) or math.isnan(wd):
print 'loss or weight norm is nan. Training Stopped!'
has_nan = True
break
average_loss += loss
if step % snapshot == 0:
saver.save(sess, logdir.snapshot_dir + '/model.ckpt', global_step=step)
sess.run([tf.local_variables_initializer()])
if step % show_period == 0:
left_hours = 0
if t0 is not None:
delta_t = (datetime.datetime.now() - t0).seconds
left_time = (max_iter - step) / show_period * delta_t
left_hours = left_time/3600.0
t0 = datetime.datetime.now()
average_loss /= show_period
if step == 0:
average_loss *= show_period
f_log.write('%d,%f,%f,%f\n' % (step, average_loss, precision, wd))
f_log.flush()
print '%s %s] Step %s, lr = %f, wd_rate = %f, wd_rate_2 = %f ' \
% (str(datetime.datetime.now()), str(os.getpid()), step, lrn_rate, wd_rate, wd_rate2)
print '\t loss = %.4f, precision = %.4f, wd = %.4f' % (average_loss, precision, wd)
print '\t estimated time left: %.1f hours. %d/%d' % (left_hours, step, max_iter)
average_loss = 0.0
coord.request_stop()
coord.join(threads)
return f_log, logdir, has_nan # f_log and logdir returned for eval.
def predict(i_ckpt):
assert i_ckpt is not None
if FLAGS.float_type == 16:
print('\n< using tf.float16 >\n')
float_type = tf.float16
else:
print('\n< using tf.float32 >\n')
float_type = tf.float32
image_size = FLAGS.test_image_size
assert FLAGS.test_image_size % 48 == 0
with tf.device('/cpu:0'):
reader = SegmentationImageReader(FLAGS.database,
FLAGS.mode, (image_size, image_size),
random_scale=False,
random_mirror=False,
random_blur=False,
random_rotate=False,
color_switch=FLAGS.color_switch)
images_pl = [tf.placeholder(tf.float32, [None, image_size, image_size, 3])]
model = pspnet_mg.PSPNetMG(reader.num_classes,
mode='val',
resnet=FLAGS.network,
data_format=FLAGS.data_format,
float_type=float_type,
has_aux_loss=False,
structure_in_paper=FLAGS.structure_in_paper)
logits = model.inference(images_pl)
probas_op = tf.nn.softmax(logits[0],
dim=1 if FLAGS.data_format == 'NCHW' else 3)
# ========================= end of building model ================================
gpu_options = tf.GPUOptions(allow_growth=False)
config = tf.ConfigProto(log_device_placement=False,
gpu_options=gpu_options,
allow_soft_placement=True)
sess = tf.Session(config=config)
sess.run(
[tf.global_variables_initializer(),
tf.local_variables_initializer()])
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
loader = tf.train.Saver(max_to_keep=0)
loader.restore(sess, i_ckpt)
print('Succesfully loaded model from %s.' % i_ckpt)
print(
'======================= eval process begins ========================='
)
if FLAGS.save_prediction == 0 and FLAGS.mode != 'test':
print('not saving prediction ... ')
average_loss = 0.0
confusion_matrix = np.zeros((reader.num_classes, reader.num_classes),
dtype=np.int64)
if FLAGS.save_prediction == 1 or FLAGS.mode == 'test':
try:
os.mkdir('./' + FLAGS.mode + '_set')
except:
pass
prefix = './' + FLAGS.mode + '_set'
try:
os.mkdir(os.path.join(prefix, FLAGS.weights_ckpt.split('/')[-2]))
except:
pass
prefix = os.path.join(prefix, FLAGS.weights_ckpt.split('/')[-2])
if FLAGS.ms == 1:
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
else:
scales = [1.0]
images_filenames = reader.image_list
labels_filenames = reader.label_list
img_mean = reader.img_mean
if FLAGS.test_max_iter is None:
max_iter = len(images_filenames)
else:
max_iter = FLAGS.test_max_iter
step = 0
while step < max_iter:
image, label = cv2.imread(images_filenames[step],
1), cv2.imread(labels_filenames[step], 0)
label = np.reshape(label, [1, label.shape[0], label.shape[1], 1])
image_height, image_width = image.shape[0], image.shape[1]
total_logits = np.zeros(
(image_height, image_width, reader.num_classes), np.float32)
for scale in scales:
imgsplitter = ImageSplitter(image, scale, FLAGS.color_switch,
image_size, img_mean)
crops = imgsplitter.get_split_crops()
# This is a suboptimal solution. More batches each iter, more rapid.
# But the limit of batch size is unknown.
# TODO: Or there should be a more efficient way.
if crops.shape[0] > 10 and FLAGS.database == 'Cityscapes':
half = crops.shape[0] // 2
feed_dict = {images_pl[0]: crops[0:half]}
[logits_0] = sess.run([probas_op], feed_dict=feed_dict)
feed_dict = {images_pl[0]: crops[half:]}
[logits_1] = sess.run([probas_op], feed_dict=feed_dict)
logits = np.concatenate((logits_0, logits_1), axis=0)
else:
feed_dict = {images_pl[0]: imgsplitter.get_split_crops()}
[logits] = sess.run([probas_op], feed_dict=feed_dict)
scale_logits = imgsplitter.reassemble_crops(logits)
if FLAGS.mirror == 1:
image_mirror = image[:, ::-1]
imgsplitter_mirror = ImageSplitter(image_mirror, scale,
FLAGS.color_switch,
image_size, img_mean)
crops_m = imgsplitter_mirror.get_split_crops()
if crops_m.shape[0] > 10:
half = crops_m.shape[0] // 2
feed_dict = {images_pl[0]: crops_m[0:half]}
[logits_0] = sess.run([probas_op], feed_dict=feed_dict)
feed_dict = {images_pl[0]: crops_m[half:]}
[logits_1] = sess.run([probas_op], feed_dict=feed_dict)
logits_m = np.concatenate((logits_0, logits_1), axis=0)
else:
feed_dict = {
images_pl[0]: imgsplitter_mirror.get_split_crops()
}
[logits_m] = sess.run([probas_op], feed_dict=feed_dict)
logits_m = imgsplitter_mirror.reassemble_crops(logits_m)
scale_logits += logits_m[:, ::-1]
if scale != 1.0:
scale_logits = cv2.resize(scale_logits,
(image_width, image_height),
interpolation=cv2.INTER_LINEAR)
total_logits += scale_logits
prediction = np.argmax(total_logits, axis=-1)
# print np.max(label), np.max(prediction)
image_prefix = images_filenames[step].split('/')[-1].split(
'.')[0] + '_' + FLAGS.weights_ckpt.split('/')[-2]
if FLAGS.database == 'Cityscapes':
cv2.imwrite(os.path.join(prefix, image_prefix + '_prediction.png'),
trainid_to_labelid(prediction))
if FLAGS.coloring == 1:
cv2.imwrite(
os.path.join(prefix, image_prefix + '_coloring.png'),
cv2.cvtColor(coloring(prediction), cv2.COLOR_BGR2RGB))
else:
cv2.imwrite(os.path.join(prefix, image_prefix + '_prediction.png'),
prediction)
# TODO: add coloring for databases other than Cityscapes.
step += 1
compute_confusion_matrix(label, prediction, confusion_matrix)
if step % 20 == 0:
print('%s %s] %d / %d. iou updating' \
% (str(datetime.datetime.now()), str(os.getpid()), step, max_iter))
compute_iou(confusion_matrix)
print(average_loss / step)
precision = compute_iou(confusion_matrix)
coord.request_stop()
coord.join(threads)
return average_loss / max_iter, precision
def eval(i_ckpt):
# does not perform multi-scale test. ms-test is in predict.py
tf.reset_default_graph()
print '================',
if FLAGS.data_type == 16:
print 'using tf.float16 ====================='
data_type = tf.float16
else:
print 'using tf.float32 ====================='
data_type = tf.float32
if 'pspnet' in FLAGS.network:
input_size = FLAGS.test_image_size
print '=====because using pspnet, the inputs have a fixed size and should be divided by 48:', input_size
assert FLAGS.test_image_size % 48 == 0
else:
input_size = None
return
with tf.device('/cpu:0'):
reader = SegmentationImageReader(
FLAGS.server,
FLAGS.database,
'val',
(input_size, input_size),
random_scale=False,
random_mirror=False,
random_blur=False,
random_rotate=False,
color_switch=FLAGS.color_switch)
images_pl = [tf.placeholder(tf.float32, [None, input_size, input_size, 3])]
labels_pl = [tf.placeholder(tf.int32, [None, input_size, input_size, 1])]
resnet = 'resnet_v1_101'
PSPModel = pspnet_mg.PSPNetMG
with tf.variable_scope(resnet):
model = PSPModel(reader.num_classes, None, None, None,
mode='val', bn_epsilon=FLAGS.epsilon, resnet=resnet,
norm_only=FLAGS.norm_only,
float_type=data_type,
has_aux_loss=False,
structure_in_paper=FLAGS.structure_in_paper,
resize_images_method=FLAGS.resize_images_method,
train_conv2dt=FLAGS.train_conv2dt
)
logits = model.inference(images_pl)
model.compute_loss(labels_pl, logits)
# ========================= end of building model ================================
gpu_options = tf.GPUOptions(allow_growth=False)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options, allow_soft_placement=True)
sess = tf.Session(config=config)
sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
if i_ckpt is not None:
loader = tf.train.Saver(max_to_keep=0)
loader.restore(sess, i_ckpt)
eval_step = i_ckpt.split('-')[-1]
print('Succesfully loaded model from %s at step=%s.' % (i_ckpt, eval_step))
print '======================= eval process begins ========================='
average_loss = 0.0
confusion_matrix = np.zeros((reader.num_classes, reader.num_classes), dtype=np.int64)
images_filenames = reader.image_list
labels_filenames = reader.label_list
img_mean = reader.img_mean
if FLAGS.test_max_iter is None:
max_iter = len(images_filenames)
else:
max_iter = FLAGS.test_max_iter
step = 0
while step < max_iter:
image, label = cv2.imread(images_filenames[step], 1), cv2.imread(labels_filenames[step], 0)
label = np.reshape(label, [1, label.shape[0], label.shape[1], 1])
imgsplitter = ImageSplitter(image, 1.0, FLAGS.color_switch, input_size, img_mean)
feed_dict = {images_pl[0]: imgsplitter.get_split_crops()}
[logits] = sess.run([
model.probabilities
],
feed_dict=feed_dict
)
total_logits = imgsplitter.reassemble_crops(logits)
if FLAGS.mirror == 1:
image_mirror = image[:, ::-1]
imgsplitter_mirror = ImageSplitter(image_mirror, 1.0, FLAGS.color_switch, input_size, img_mean)
feed_dict = {images_pl[0]: imgsplitter_mirror.get_split_crops()}
[logits_m] = sess.run([
model.probabilities
],
feed_dict=feed_dict
)
logits_m = imgsplitter_mirror.reassemble_crops(logits_m)
total_logits += logits_m[:, ::-1]
prediction = np.argmax(total_logits, axis=-1)
step += 1
compute_confusion_matrix(label, prediction, confusion_matrix)
if step % 20 == 0:
print '%s %s] %d / %d. iou updating' \
% (str(datetime.datetime.now()), str(os.getpid()), step, max_iter)
compute_iou(confusion_matrix)
print 'imprecise loss', average_loss / step
precision = compute_iou(confusion_matrix)
coord.request_stop()
coord.join(threads)
return average_loss / max_iter, precision
def predict(i_ckpt):
tf.reset_default_graph()
print '================',
if FLAGS.data_type == 16:
print 'using tf.float16 ====================='
data_type = tf.float16
else:
print 'using tf.float32 ====================='
data_type = tf.float32
image_size = FLAGS.test_image_size
print '=====because using pspnet, the inputs have a fixed size and should be divided by 48:', image_size
assert FLAGS.test_image_size % 48 == 0
with tf.device('/cpu:0'):
reader = SegmentationImageReader(FLAGS.server,
FLAGS.database,
FLAGS.mode, (image_size, image_size),
random_scale=False,
random_mirror=False,
random_blur=False,
random_rotate=False,
color_switch=FLAGS.color_switch)
images_pl = [tf.placeholder(tf.float32, [None, image_size, image_size, 3])]
labels_pl = [tf.placeholder(tf.int32, [None, image_size, image_size, 1])]
with tf.variable_scope('resnet_v1_101'):
model = pspnet_mg.PSPNetMG(
reader.num_classes,
None,
None,
None,
mode='val',
bn_epsilon=FLAGS.epsilon,
resnet='resnet_v1_101',
norm_only=FLAGS.norm_only,
float_type=data_type,
has_aux_loss=False,
structure_in_paper=FLAGS.structure_in_paper,
resize_images_method=FLAGS.resize_images_method)
l = model.inference(images_pl)
# ========================= end of building model ================================
gpu_options = tf.GPUOptions(allow_growth=False)
config = tf.ConfigProto(log_device_placement=False,
gpu_options=gpu_options,
allow_soft_placement=True)
sess = tf.Session(config=config)
sess.run(
[tf.global_variables_initializer(),
tf.local_variables_initializer()])
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
if i_ckpt is not None:
loader = tf.train.Saver(max_to_keep=0)
loader.restore(sess, i_ckpt)
eval_step = i_ckpt.split('-')[-1]
print('Succesfully loaded model from %s at step=%s.' %
(i_ckpt, eval_step))
print '======================= eval process begins ========================='
if FLAGS.save_prediction == 0 and FLAGS.mode != 'test':
print 'not saving prediction ... '
average_loss = 0.0
confusion_matrix = np.zeros((reader.num_classes, reader.num_classes),
dtype=np.int64)
if FLAGS.save_prediction == 1 or FLAGS.mode == 'test':
try:
os.mkdir('./' + FLAGS.mode + '_set')
except:
pass
prefix = './' + FLAGS.mode + '_set'
try:
os.mkdir(os.path.join(prefix, FLAGS.weights_ckpt.split('/')[-2]))
except:
pass
prefix = os.path.join(prefix, FLAGS.weights_ckpt.split('/')[-2])
if FLAGS.ms == 1:
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
else:
scales = [1.0]
images_filenames = reader.image_list
labels_filenames = reader.label_list
if FLAGS.test_max_iter is None:
max_iter = len(images_filenames)
else:
max_iter = FLAGS.test_max_iter
# IMG_MEAN = [123.680000305, 116.778999329, 103.939002991] # RGB mean from official PSPNet
step = 0
while step < max_iter:
image, label = cv2.imread(images_filenames[step],
1), cv2.imread(labels_filenames[step], 0)
label = np.reshape(label, [1, label.shape[0], label.shape[1], 1])
image_height, image_width = image.shape[0], image.shape[1]
total_logits = np.zeros(
(image_height, image_width, reader.num_classes), np.float32)
for scale in scales:
imgsplitter = ImageSplitter(image, scale, FLAGS.color_switch,
image_size, reader.img_mean)
crops = imgsplitter.get_split_crops()
# This is a suboptimal solution. More batches each iter, more rapid.
# But the limit of batch size is unknown.
# TODO: Or there should be a more efficient way.
if crops.shape[0] > 10:
half = crops.shape[0] / 2
feed_dict = {images_pl[0]: crops[0:half]}
[logits_0] = sess.run([model.probabilities],
feed_dict=feed_dict)
feed_dict = {images_pl[0]: crops[half:]}
[logits_1] = sess.run([model.probabilities],
feed_dict=feed_dict)
logits = np.concatenate((logits_0, logits_1), axis=0)
else:
feed_dict = {images_pl[0]: imgsplitter.get_split_crops()}
[logits] = sess.run([model.probabilities], feed_dict=feed_dict)
scale_logits = imgsplitter.reassemble_crops(logits)
if FLAGS.mirror == 1:
image_mirror = image[:, ::-1]
imgsplitter_mirror = ImageSplitter(image_mirror, scale,
FLAGS.color_switch,
image_size, reader.img_mean)
crops_m = imgsplitter_mirror.get_split_crops()
if crops_m.shape[0] > 10:
half = crops_m.shape[0] / 2
feed_dict = {images_pl[0]: crops_m[0:half]}
[logits_0] = sess.run([model.probabilities],
feed_dict=feed_dict)
feed_dict = {images_pl[0]: crops_m[half:]}
[logits_1] = sess.run([model.probabilities],
feed_dict=feed_dict)
logits_m = np.concatenate((logits_0, logits_1), axis=0)
else:
feed_dict = {
images_pl[0]: imgsplitter_mirror.get_split_crops()
}
[logits_m] = sess.run([model.probabilities],
feed_dict=feed_dict)
logits_m = imgsplitter_mirror.reassemble_crops(logits_m)
scale_logits += logits_m[:, ::-1]
if scale != 1.0:
scale_logits = cv2.resize(scale_logits,
(image_width, image_height),
interpolation=cv2.INTER_LINEAR)
total_logits += scale_logits
prediction = np.argmax(total_logits, axis=-1)
# print np.max(label), np.max(prediction)
if FLAGS.database == 'Cityscapes' and (FLAGS.save_prediction == 1
or FLAGS.mode == 'test'):
image_prefix = images_filenames[step].split('/')[-1].split('_leftImg8bit.png')[0] + '_' \
+ FLAGS.weights_ckpt.split('/')[-2]
cv2.imwrite(os.path.join(prefix, image_prefix + '_prediction.png'),
trainid_to_labelid(prediction))
if FLAGS.coloring == 1:
color_prediction = coloring(prediction)
cv2.imwrite(
os.path.join(prefix, image_prefix + '_coloring.png'),
cv2.cvtColor(color_prediction, cv2.COLOR_BGR2RGB))
elif FLAGS.database == 'SBD' and (FLAGS.save_prediction == 1
or FLAGS.mode == 'test'):
image_prefix = images_filenames[step].split('/')[-1].split(
'.jpg')[0]
cv2.imwrite(os.path.join(prefix, image_prefix + '.png'),
prediction)
else:
pass
step += 1
compute_confusion_matrix(label, prediction, confusion_matrix)
if step % 20 == 0:
print '%s %s] %d / %d. iou updating' \
% (str(datetime.datetime.now()), str(os.getpid()), step, max_iter)
compute_iou(confusion_matrix)
print average_loss / step
precision = compute_iou(confusion_matrix)
coord.request_stop()
coord.join(threads)
return average_loss / max_iter, precision