def run(self, train_reader=None, valid_reader=None, test_reader=None, pretrained_model_dir=None, layers2load=None,
isTrain=False, img_mean=np.array((0, 0, 0), dtype=np.float32), verb_step=100, save_epoch=5, gpu=None,
tile_size=(5000, 5000), patch_size=(572, 572), truth_val=1, continue_dir=None, load_epoch_num=None,
fineTune=False, valid_iou=False, best_model=True):
if gpu is not None:
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
if isTrain:
coord = tf.train.Coordinator()
with tf.Session(config=self.config) as sess:
# init model
init = [tf.global_variables_initializer(), tf.local_variables_initializer()]
sess.run(init)
saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=1)
# load model
if pretrained_model_dir is not None:
if layers2load is not None:
self.load_weights(pretrained_model_dir, layers2load)
else:
if not fineTune:
restore_var = [v for v in tf.global_variables() if 'resnet_v1' in v.name and
not 'Adam' in v.name]
loader = tf.train.Saver(var_list=restore_var)
self.load(pretrained_model_dir, sess, loader, epoch=load_epoch_num)
else:
self.load(pretrained_model_dir, sess, saver, epoch=load_epoch_num)
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
try:
train_summary_writer = tf.summary.FileWriter(self.ckdir, sess.graph)
self.train('X', 'Y', self.n_train, sess, train_summary_writer,
n_valid=self.n_valid, train_reader=train_reader, valid_reader=valid_reader,
image_summary=util_functions.image_summary, img_mean=img_mean,
verb_step=verb_step, save_epoch=save_epoch, continue_dir=continue_dir,
valid_iou=valid_iou)
finally:
coord.request_stop()
coord.join(threads)
saver.save(sess, '{}/model.ckpt'.format(self.ckdir), global_step=self.global_step)
else:
pad = self.crop_pixel * 2 + self.get_overlap()
with tf.Session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
self.load(pretrained_model_dir, sess, epoch=load_epoch_num, best_model=best_model)
self.model_name = pretrained_model_dir.split('/')[-1]
result = self.test('X', sess, test_reader)
image_pred = uabUtilreader.un_patchify_shrink(result,
[tile_size[0] + pad, tile_size[1] + pad],
tile_size,
patch_size,
[patch_size[0] - pad, patch_size[1] - pad],
overlap=pad)
return util_functions.get_pred_labels(image_pred) * truth_val
overlap=model.get_overlap(),
padding=np.array((model.get_overlap() / 2, model.get_overlap() / 2)),
isTrain=False)
test_reader = reader.readManager
# run algo
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
with tf.Session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
model.load(sp_model_dir, sess)
result = model.test('X', sess, test_reader)
image_pred = uabUtilreader.un_patchify_shrink(result,
[tile_size[0] + model.get_overlap(),
tile_size[1] + model.get_overlap()],
tile_size, input_size,
[input_size[0] - model.get_overlap(),
input_size[1] - model.get_overlap()],
overlap=model.get_overlap())
pred = util_functions.get_pred_labels(image_pred)
mask_img = np.copy(orig_img)
mask_img = draw_contour(mask_img, pred, truth_val=1)
# mask_img = util_functions.add_mask(mask_img, pred, [255, None, None], mask_1=1)
region = [250, 450, 450, 650]
save_dir = r'/media/ei-edl01/user/bh163/figs/2018.03.15.gbdx/blog_figures'
plt.figure(figsize=(8, 8))
plt.imshow(mask_img[region[0]:region[1], region[2]:region[3], :])
plt.tight_layout()
plt.show()
padding=np.array(
(model.get_overlap() // 2,
model.get_overlap() // 2)),
isTrain=False,
shift=slide_step)
rManager = reader.readManager
with tf.Session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
model.load(model_dir, sess, epoch=None, best_model=True)
model.model_name = model_dir.split('/')[-1]
result = model.test('X', sess, rManager)
pad = model.get_overlap()
image_pred = uabUtilreader.un_patchify_shrink(
result, [tile_size[0] + pad, 4576 + pad], [5000, 4576],
input_size, [input_size[0] - pad, input_size[1] - pad],
overlap=pad)
pred_overall = util_functions.get_pred_labels(image_pred) * 1
pred_overall = np.roll(pred_overall, shift=slide_step, axis=1)
pred_overall = pred_overall[:, 1000:-1000]
#pred_overall = pred_overall[:, shift_max-slide_step:-slide_step-1]
truth_label_img = imageio.imread(
os.path.join(parent_dir_truth, file_name_truth))
#truth_label_img = np.roll(truth_label_img, -slide_step, axis=1)
truth_label_img = truth_label_img[:, :4576]
truth_label_img = truth_label_img[:, 1000:-1000]
iou = util_functions.iou_metric(truth_label_img,
pred_overall,
divide_flag=True)
duration = time.time() - start_time
if iou[1] != 0:
chipFiles=[file_name],
chip_size=input_size,
tile_size=tile_size,
batchSize=batch_size,
block_mean=img_mean,
overlap=model.get_overlap(),
padding=pad,
isTrain=False)
rManager = reader.readManager
# run the model
with tf.Session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
model.load(model_dir, sess)
model.model_name = model_dir.split('/')[-1]
result = model.test('X', sess, rManager)
image_pred = uabUtilreader.un_patchify_shrink(
result, [tile_size[0] + 184, tile_size[1] + 184],
tile_size,
input_size, [input_size[0] - 184, input_size[1] - 184],
overlap=92)
tf.reset_default_graph()
run_time = time.time() - start_time
# save run time
print(run_time)
with open(save_file_name, 'a') as f:
f.write('{} {}\n'.format(size, run_time))
