def evaluation(config_file):
config = parse_config(config_file)['evaluation']
metric = config['metric']
labels = config['label_list']
label_convert_source = config.get('label_convert_source', None)
label_convert_target = config.get('label_convert_target', None)
s_folder = config['segmentation_folder']
g_folder = config['ground_truth_folder']
s_postfix = config.get('segmentation_postfix',None)
g_postfix = config.get('ground_truth_postfix',None)
s_format = config.get('segmentation_format', "nii.gz")
g_format = config.get('ground_truth_format', "nii.gz")
remove_outlier = config.get('remove_outlier', False)
s_postfix_long = '.' + s_format
if(s_postfix is not None):
s_postfix_long = '_' + s_postfix + s_postfix_long
g_postfix_long = '.' + g_format
if(g_postfix is not None):
g_postfix_long = '_' + g_postfix + g_postfix_long
patient_names_file = config['patient_file_names']
with open(patient_names_file) as f:
content = f.readlines()
patient_names = [x.strip() for x in content]
score_all_data = []
for i in range(len(patient_names)):
s_name = os.path.join(s_folder, patient_names[i] + s_postfix_long)
g_name = os.path.join(g_folder, patient_names[i] + g_postfix_long)
s_volume, s_spacing = load_image_as_array(s_name, with_spacing = True)
g_volume, g_spacing = load_image_as_array(g_name, with_spacing = True)
if((label_convert_source is not None) and label_convert_target is not None):
s_volume = convert_label(s_volume, label_convert_source, label_convert_target)
# fuse multiple labels
s_volume_sub = np.zeros_like(s_volume)
g_volume_sub = np.zeros_like(g_volume)
for lab in labels:
s_volume_sub = s_volume_sub + np.asarray(s_volume == lab, np.uint8)
g_volume_sub = g_volume_sub + np.asarray(g_volume == lab, np.uint8)
# if(s_volume_sub.sum() > 0):
# s_volume_sub = get_largest_component(s_volume_sub)
if(remove_outlier):
strt = ndimage.generate_binary_structure(3,2) # iterate structure
post = ndimage.morphology.binary_closing(s_volume_sub, strt)
post = get_largest_component(post)
s_volume_sub = np.asarray(post*s_volume_sub, np.uint8)
temp_score = get_evaluation_score(s_volume_sub > 0, g_volume_sub > 0,
g_spacing, metric)
score_all_data.append(temp_score)
print(patient_names[i], temp_score)
score_all_data = np.asarray(score_all_data)
score_mean = [score_all_data.mean(axis = 0)]
score_std = [score_all_data.std(axis = 0)]
np.savetxt("{0:}/{1:}_all_temp.txt".format(s_folder, metric), score_all_data)
np.savetxt("{0:}/{1:}_mean_temp.txt".format(s_folder, metric), score_mean)
np.savetxt("{0:}/{1:}_std_temp.txt".format(s_folder, metric), score_std)
print("{0:} mean ".format(metric), score_mean)
print("{0:} std ".format(metric), score_std)
def model_train(config_file):
config = parse_config(config_file)
app_type = config['training']['app_type']
if (app_type == 0):
train_agent = SegmentationTrainAgent(config)
else:
train_agent = RegressionTrainAgent(config)
train_agent.construct_network()
train_agent.create_optimization_step_and_data_generator()
train_agent.train()
def post_process(config_file):
config = parse_config(config_file)['post_process']
in_folder = config['input_folder']
out_folder = config['output_folder']
s_postfix = config.get('segmentation_postfix',None)
s_postfix = '.nii.gz' if (s_postfix is None) else '_' + s_postfix + '.nii.gz'
file_names = os.listdir(in_folder)
file_names = [name for name in file_names if 'nii.gz' in name]
for file_name in file_names:
in_name = os.path.join(in_folder, file_name)
out_name = os.path.join(out_folder, file_name)
seg = load_nifty_volume_as_array(in_name)
strt = ndimage.generate_binary_structure(3,2) # iterate structure
post = ndimage.morphology.binary_closing(seg, strt)
post = get_largest_component(post)
post = np.asarray(post*seg, np.uint8)
save_array_as_nifty_volume(post, out_name)
def test_generator(config_file):
config = parse_config(config_file)
config_data = config['dataset']
temp_dir = config_data['temp_dir']
batch_size = config['sampler']['batch_size']
# Place data loading and preprocessing on the cpu
tf.set_random_seed(0)
with tf.device('/cpu:0'):
tr_data = ImageDataGenerator(config_data['data_train'],
config['sampler'])
# create an reinitializable iterator given the dataset structure
iterator = Iterator.from_structure(tr_data.data.output_types,
tr_data.data.output_shapes)
next_batch = iterator.get_next()
# Ops for initializing the two different iterators
training_init_op = iterator.make_initializer(tr_data.data)
#validation_init_op = iterator.make_initializer(val_data.data)
train_batches_per_epoch = config['training']['batch_number']
num_epochs = config['training']['maximal_epoch']
app_type = config['training']['app_type']
# Start Tensorflow session
with tf.Session() as sess:
# Initialize all variables
sess.run(tf.global_variables_initializer())
total_step = 0
# Loop over number of epochs
for epoch in range(num_epochs):
print("{} Epoch number: {}".format(datetime.now(), epoch + 1))
# Initialize iterator with the training dataset
sess.run(training_init_op)
for step in range(train_batches_per_epoch):
if (app_type == 0):
[img_batch, weight_batch,
label_batch] = sess.run(next_batch)
img_0 = img_batch[0, :, :, :, :]
lab_0 = label_batch[0, :, :, :, 0]
print(epoch, step, img_0.shape, lab_0.shape)
if (config['training']['save_postfix'] == 'nii.gz'):
img_0 = img_0[:, :, :, 0]
img_d = img_0.shape[0]
lab_d = lab_0.shape[0]
if (lab_d < img_d):
margin = (img_d - lab_d) / 2
pad_lab = np.zeros_like(img_0)
pad_lab[np.ix_(range(margin, margin + lab_d),
range(lab_0.shape[1]),
range(lab_0.shape[2]))] = lab_0
lab_0 = pad_lab
save_array_as_nifty_volume(
img_0,
'{0:}/img{1:}.nii.gz'.format(temp_dir, total_step))
save_array_as_nifty_volume(
lab_0,
'{0:}/lab{1:}.nii.gz'.format(temp_dir, total_step))
elif (config['training']['save_postfix'] == 'jpg'):
iten_mean = np.asarray(
[179.69427237, 146.44891944, 134.39686832])
iten_std = np.asarray(
[40.37515566, 42.92464467, 46.74197245])
img_0 = img_0[0] * iten_std + iten_mean
img_0 = np.asarray(img_0, np.uint8)
lab_0 = np.asarray(lab_0[0], np.uint8) * 255
print(img_0.min(), img_0.max(), lab_0.min(),
lab_0.max())
img_0 = Image.fromarray(img_0, 'RGB')
lab_0 = Image.fromarray(lab_0, 'L')
img_0.save('{0:}/img{1:}.jpg'.format(
temp_dir, total_step))
lab_0.save('{0:}/lab{1:}.jpg'.format(
temp_dir, total_step))
else:
[img_batch, stp] = sess.run(next_batch)
print(stp)
img_0 = img_batch[0, 0, :, :, 0]
plt.imshow(img_0)
plt.show()
total_step = total_step + 1
def convert_to_rf_records(config_file):
config = parse_config(config_file)
config_data = config['data']
data_loader = DataLoader(config_data)
data_loader.load_data()
data_loader.save_to_tfrecords()
if __name__ == '__main__':
if (('--input_names' not in sys.argv)
or ('--segment_output_names' not in sys.argv)):
print('Inccorrect command line, please follow the format:')
print(' python fetal_brain_seg.py ' +
' --input_names img1.nii.gz img2.nii.gz ...' +
' --segment_output_names img1_seg.nii.gz img2_seg.nii.gz ...' +
' (--detect_output_names img1_det.nii.gz img2_det.nii.gz ...')
exit()
idx_0 = sys.argv.index('--input_names')
idx_1 = sys.argv.index('--segment_output_names')
input_names = sys.argv[idx_0 + 1:idx_1]
img_num = idx_1 - idx_0 - 1
seg_names = sys.argv[idx_1 + 1:idx_1 + 1 + img_num]
det_names = None
if ('--detect_output_names' in sys.argv):
idx_2 = sys.argv.index('--detect_output_names')
det_names = sys.argv[idx_2 + 1:idx_2 + 1 + img_num]
data_config = {}
for i in range(len(input_names)):
img_i = 'image_{0:}'.format(i)
data_config[img_i] = {}
data_config[img_i]['input'] = input_names[i]
data_config[img_i]['segment_output'] = seg_names[i]
if (det_names is not None):
data_config[img_i]['detect_output'] = det_names[i]
net_config_file = 'cfg_net.txt'
net_config = parse_config(net_config_file)
model_test(net_config, data_config)
save_name = "{0:}_{1:}.{2:}".format(patient_name,
self.config_data['output_postfix'],
self.config_data['outputfile_postfix'])
if (self.config_data['outputfile_postfix'] == "nii.gz" or self.config_data['outputfile_postfix'] == "nii"):
save_array_as_nifty_volume(out, self.config_data['save_root']+'/'+save_name, file_names[0])
elif(self.config_data['outputfile_postfix'] == "jpg" or self.config_data['outputfile_postfix'] == "png"):
assert(out.shape[0] == 1 and len(out.shape) == 3)
out = np.reshape(out, out.shape[1:])
if(self.config_net['class_num'] == 2):
out = out * 255
out_img = Image.fromarray(out, 'L')
out_img.save(self.config_data['save_root'] + '/' + save_name)
print("save array with shape", out.shape)
if(self.config_test.get('save_probability', False)):
save_name = '{0:}_{1:}.nii.gz'.format(patient_name, 'Prob')
save_array_as_nifty_volume(outp, self.config_data['save_root']+'/'+save_name, file_names[0])
test_time = np.asarray(test_time)
print('test time', test_time.mean(), test_time.std())
np.savetxt("{0:}/test_time.txt".format(self.config_data['save_root']), test_time)
if __name__ == '__main__':
if(len(sys.argv) != 2):
print('Number of arguments should be 2. e.g.')
print(' python train_test/model_test.py config.txt')
exit()
config_file = str(sys.argv[1])
assert(os.path.isfile(config_file))
config = parse_config(config_file)
test_agent = TestAgent(config)
test_agent.test()
# stage 2, segment
margin = [0, 10, 10]
bb_min, bb_max = get_ND_bounding_box(detect_out, margin)
img_roi = crop_ND_volume_with_bounding_box(img, bb_min + [0],
bb_max + [0])
outp_roi = segment_agent.test_one_volume(img_roi)
out_roi = np.asarray(outp_roi > 0.5, np.uint8)
strt = ndimage.generate_binary_structure(3, 2)
post = padded_binary_closing(out_roi, strt)
post = get_largest_component(post)
out_roi = out_roi * post
out = np.zeros(img.shape[:-1], np.uint8)
out = set_ND_volume_roi_with_bounding_box_range(
out, bb_min, bb_max, out_roi)
save_array_as_nifty_volume(out, segment_name, input_name)
if __name__ == '__main__':
if (len(sys.argv) < 2):
print('Number of arguments should be 2. e.g.')
print(' python test.py cfg_data.txt')
exit()
net_config_file = 'cfg_net.txt'
data_config_file = sys.argv[1]
assert (os.path.isfile(net_config_file)
and os.path.isfile(data_config_file))
net_config = parse_config(net_config_file)
data_config = parse_config(data_config_file)
model_test(net_config, data_config)
def model_train(config_file):
config = parse_config(config_file)
train_agent = CustomTrainTestAgent(config, 'train')
train_agent.train()