def init_networks(self):
"""
Initialize networks and placeholders.
"""
network_image_size = list(reversed(self.image_size))
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([('image', [1] + network_image_size),
('single_label', [self.num_labels] + network_image_size),
('single_heatmap', [1] + network_image_size)])
else:
data_generator_entries = OrderedDict([('image', network_image_size + [1]),
('single_label', network_image_size + [self.num_labels]),
('single_heatmap', network_image_size + [1])])
data_generator_types = {'image': tf.float32,
'labels': tf.uint8}
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(coord=self.coord, dataset=self.dataset_train, data_names_and_shapes=data_generator_entries, data_types=data_generator_types, batch_size=self.batch_size)
data, mask, single_heatmap = self.train_queue.dequeue()
data_heatmap_concat = tf.concat([data, single_heatmap], axis=1)
prediction = training_net(data_heatmap_concat, num_labels=self.num_labels, is_training=True, actual_network=self.unet, padding=self.padding, **self.network_parameters)
# losses
self.loss_net = self.loss_function(labels=mask, logits=prediction, data_format=self.data_format)
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss = self.loss_net + self.loss_reg
# solver
global_step = tf.Variable(self.current_iter, trainable=False)
learning_rate = tf.train.piecewise_constant(global_step, self.learning_rate_boundaries, self.learning_rates)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
unclipped_gradients, variables = zip(*optimizer.compute_gradients(self.loss))
norm = tf.global_norm(unclipped_gradients)
if self.clip_gradient_global_norm > 0:
gradients, _ = tf.clip_by_global_norm(unclipped_gradients, self.clip_gradient_global_norm)
else:
gradients = unclipped_gradients
self.optimizer = optimizer.apply_gradients(zip(gradients, variables), global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net), ('loss_reg', self.loss_reg), ('gradient_norm', norm)])
# build val graph
self.data_val, self.mask_val, self.single_heatmap_val = create_placeholders_tuple(data_generator_entries, data_types=data_generator_types, shape_prefix=[1])
self.data_heatmap_concat_val = tf.concat([self.data_val, self.single_heatmap_val], axis=1)
self.prediction_val = training_net(self.data_heatmap_concat_val, num_labels=self.num_labels, is_training=False, actual_network=self.unet, padding=self.padding, **self.network_parameters)
self.prediction_softmax_val = tf.nn.sigmoid(self.prediction_val)
if self.has_validation_groundtruth:
self.loss_val = self.loss_function(labels=self.mask_val, logits=self.prediction_val, data_format=self.data_format)
self.val_losses = OrderedDict([('loss', self.loss_val), ('loss_reg', self.loss_reg), ('gradient_norm', tf.constant(0, tf.float32))])
def initNetworks(self):
net = tf.make_template('net', self.network)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([('image', [self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks, 3])])
data_generator_entries_val = OrderedDict([('image', [self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks, 3])])
else:
data_generator_entries = OrderedDict([('image', list(reversed(self.image_size)) + [self.image_channels]),
('landmarks', [self.num_landmarks, 3])])
data_generator_entries_val = OrderedDict([('image', list(reversed(self.image_size)) + [self.image_channels]),
('landmarks', [self.num_landmarks, 3])])
sigmas = tf.get_variable('sigmas', [self.num_landmarks], initializer=tf.constant_initializer(self.heatmap_sigma))
sigmas_list = [(f's{i}', sigmas[i]) for i in range(self.num_landmarks)]
# build training graph
self.train_queue = DataGenerator(self.dataset_train, self.coord, data_generator_entries, batch_size=self.batch_size, n_threads=8)
placeholders = self.train_queue.dequeue()
image = placeholders[0]
target_landmarks = placeholders[1]
prediction = net(image, num_landmarks=self.num_landmarks, is_training=True, data_format=self.data_format)
target_heatmaps = generate_heatmap_target(list(reversed(self.heatmap_size)), target_landmarks, sigmas, scale=self.sigma_scale, normalize=True, data_format=self.data_format)
loss_sigmas = self.loss_sigmas(sigmas, target_landmarks)
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss_net = self.loss_function(target_heatmaps, prediction)
self.loss = self.loss_net + tf.cast(self.loss_reg, tf.float32) + loss_sigmas
# optimizer
global_step = tf.Variable(self.current_iter, trainable=False)
optimizer = tf.train.MomentumOptimizer(learning_rate=self.learning_rate, momentum=0.99, use_nesterov=True)
unclipped_gradients, variables = zip(*optimizer.compute_gradients(self.loss))
norm = tf.global_norm(unclipped_gradients)
gradients, _ = tf.clip_by_global_norm(unclipped_gradients, 10000.0)
self.optimizer = optimizer.apply_gradients(zip(gradients, variables), global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net), ('loss_reg', self.loss_reg), ('loss_sigmas', loss_sigmas), ('norm', norm)] + sigmas_list)
# build val graph
self.val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(data_generator_entries_val, shape_prefix=[1])
self.image_val = self.val_placeholders['image']
self.target_landmarks_val = self.val_placeholders['landmarks']
self.prediction_val = net(self.image_val, num_landmarks=self.num_landmarks, is_training=False, data_format=self.data_format)
self.target_heatmaps_val = generate_heatmap_target(list(reversed(self.heatmap_size)), self.target_landmarks_val, sigmas, scale=self.sigma_scale, normalize=True, data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.target_heatmaps_val, self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val), ('loss_reg', self.loss_reg), ('loss_sigmas', tf.constant(0, tf.float32)), ('norm', tf.constant(0, tf.float32))] + sigmas_list)
def init_networks(self):
"""
Initialize networks and placeholders.
"""
network_image_size = list(reversed(self.image_size))
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([('image', [1] + network_image_size),
('spine_heatmap', [1] + network_image_size)])
else:
data_generator_entries = OrderedDict([('image', network_image_size + [1]),
('spine_heatmap', [1] + network_image_size)])
data_generator_types = {'image': tf.float32,
'spine_heatmap': tf.float32}
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(coord=self.coord, dataset=self.dataset_train, data_names_and_shapes=data_generator_entries, data_types=data_generator_types, batch_size=self.batch_size)
data, target_spine_heatmap = self.train_queue.dequeue()
prediction = training_net(data, num_labels=self.num_labels, is_training=True, actual_network=self.unet, padding=self.padding, **self.network_parameters)
self.loss_net = self.loss_function(target=target_spine_heatmap, pred=prediction)
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss = self.loss_net + tf.cast(self.loss_reg, tf.float32)
# solver
global_step = tf.Variable(self.current_iter, trainable=False)
learning_rate = tf.train.piecewise_constant(global_step, self.learning_rate_boundaries, self.learning_rates)
self.optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(self.loss, global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net), ('loss_reg', self.loss_reg)])
# build val graph
self.data_val, self.target_spine_heatmap_val = create_placeholders_tuple(data_generator_entries, data_types=data_generator_types, shape_prefix=[1])
self.prediction_val = training_net(self.data_val, num_labels=self.num_labels, is_training=False, actual_network=self.unet, padding=self.padding, **self.network_parameters)
if self.has_validation_groundtruth:
self.loss_val = self.loss_function(target=self.target_spine_heatmap_val, pred=self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val), ('loss_reg', self.loss_reg)])
def initNetworks(self):
net = tf.make_template('net', self.network)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([('image', [self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks] + list(reversed(self.heatmap_size)))])
data_generator_entries_val = OrderedDict([('image', [self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks] + list(reversed(self.heatmap_size)))])
else:
raise NotImplementedError
self.train_queue = DataGenerator(self.dataset_train, self.coord, data_generator_entries, batch_size=self.batch_size, n_threads=8)
placeholders = self.train_queue.dequeue()
image = placeholders[0]
landmarks = placeholders[1]
prediction = net(image, num_landmarks=self.num_landmarks, is_training=True, data_format=self.data_format)
self.loss_net = self.loss_function(landmarks, prediction)
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
if self.reg_constant > 0:
reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg
else:
self.loss_reg = 0
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net), ('loss_reg', self.loss_reg)])
#self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss)
self.optimizer = tf.train.MomentumOptimizer(learning_rate=self.learning_rate, momentum=0.99, use_nesterov=True).minimize(self.loss)
# build val graph
self.val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(data_generator_entries_val, shape_prefix=[1])
self.image_val = self.val_placeholders['image']
self.landmarks_val = self.val_placeholders['landmarks']
self.prediction_val = net(self.image_val, num_landmarks=self.num_landmarks, is_training=False, data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.landmarks_val, self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val), ('loss_reg', self.loss_reg)])
def init_networks(self):
"""
Initialize networks and placeholders.
"""
network_image_size = list(reversed(self.image_size))
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image', [1] + network_image_size),
('landmarks', [self.num_landmarks, 4]),
('landmark_mask', [1] + network_image_size)
])
else:
data_generator_entries = OrderedDict([
('image', network_image_size + [1]),
('landmarks', [self.num_landmarks, 4]),
('landmark_mask', network_image_size + [1])
])
data_generator_types = {'image': tf.float32}
# create sigmas variable
sigmas = tf.get_variable('sigmas', [self.num_landmarks],
initializer=tf.constant_initializer(
self.heatmap_sigma))
if not self.learnable_sigma:
sigmas = tf.stop_gradient(sigmas)
mean_sigmas = tf.reduce_mean(sigmas)
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(
coord=self.coord,
dataset=self.dataset_train,
data_names_and_shapes=data_generator_entries,
data_types=data_generator_types,
batch_size=self.batch_size)
data, target_landmarks, landmark_mask = self.train_queue.dequeue()
target_heatmaps = generate_heatmap_target(list(
reversed(self.heatmap_size)),
target_landmarks,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
prediction, local_prediction, spatial_prediction = training_net(
data,
num_labels=self.num_landmarks,
is_training=True,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
# losses
self.loss_net = self.loss_function(target=target_heatmaps,
pred=prediction,
mask=landmark_mask)
self.loss_sigmas = self.loss_function_sigmas(sigmas,
target_landmarks[0, :, 0])
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss = self.loss_net + self.loss_reg + self.loss_sigmas
# solver
global_step = tf.Variable(self.current_iter, trainable=False)
learning_rate = tf.train.piecewise_constant(
global_step, self.learning_rate_boundaries, self.learning_rates)
optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate,
momentum=0.99,
use_nesterov=True)
unclipped_gradients, variables = zip(
*optimizer.compute_gradients(self.loss))
norm = tf.global_norm(unclipped_gradients)
if self.clip_gradient_global_norm > 0:
gradients, _ = tf.clip_by_global_norm(
unclipped_gradients, self.clip_gradient_global_norm)
else:
gradients = unclipped_gradients
self.optimizer = optimizer.apply_gradients(zip(gradients, variables),
global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg),
('loss_sigmas', self.loss_sigmas),
('mean_sigmas', mean_sigmas),
('gradient_norm', norm)])
# build val graph
self.data_val, self.target_landmarks_val, self.landmark_mask_val = create_placeholders_tuple(
data_generator_entries,
data_types=data_generator_types,
shape_prefix=[1])
self.target_heatmaps_val = generate_heatmap_target(
list(reversed(self.heatmap_size)),
self.target_landmarks_val,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
self.prediction_val, self.local_prediction_val, self.spatial_prediction_val = training_net(
self.data_val,
num_labels=self.num_landmarks,
is_training=False,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
if self.has_validation_groundtruth:
self.loss_val = self.loss_function(target=self.target_heatmaps_val,
pred=self.prediction_val)
self.val_losses = OrderedDict([
('loss', self.loss_val), ('loss_reg', self.loss_reg),
('loss_sigmas', tf.constant(0, tf.float32)),
('mean_sigmas', tf.constant(0, tf.float32)),
('gradient_norm', tf.constant(0, tf.float32))
])
class MainLoop(MainLoopBase):
def __init__(self, cv, network_id):
super().__init__()
self.cv = cv
self.network_id = network_id
self.output_folder = network_id
if cv != -1:
self.output_folder += '_cv{}'.format(cv)
self.output_folder += '/' + self.output_folder_timestamp()
self.batch_size = 1
learning_rates = {'scn': 0.00000005, 'unet': 0.000000005}
max_iters = {'scn': 40000, 'unet': 80000}
self.learning_rate = learning_rates[self.network_id]
self.max_iter = max_iters[self.network_id]
self.test_iter = 2500
self.disp_iter = 100
self.snapshot_iter = self.test_iter
self.test_initialization = False
self.current_iter = 0
self.reg_constant = 0.0005
self.sigma_regularization = 100
self.sigma_scale = 1000
self.invert_transformation = False
self.num_landmarks = 26
self.image_size = [96, 96, 192]
self.image_spacing = [2, 2, 2]
self.heatmap_size = self.image_size
self.image_channels = 1
self.heatmap_sigma = 4
self.data_format = 'channels_first'
self.save_debug_images = False
self.base_folder = 'spine_localization_dataset'
self.generate_landmarks = True
self.cropped_training = True
self.cropped_inc = [0, 64, 0, 0]
if self.cropped_training:
dataset = Dataset(self.image_size,
self.image_spacing,
self.heatmap_sigma,
self.num_landmarks,
self.base_folder,
self.cv,
self.data_format,
self.save_debug_images,
generate_heatmaps=not self.generate_landmarks,
generate_landmarks=self.generate_landmarks)
self.dataset_train = dataset.dataset_train()
dataset = Dataset(self.image_size,
self.image_spacing,
self.heatmap_sigma,
self.num_landmarks,
self.base_folder,
self.cv,
self.data_format,
self.save_debug_images,
generate_heatmaps=not self.generate_landmarks,
generate_landmarks=self.generate_landmarks)
self.dataset_val = dataset.dataset_val()
else:
dataset = Dataset(self.image_size,
self.image_spacing,
self.heatmap_sigma,
self.num_landmarks,
self.base_folder,
self.cv,
self.data_format,
self.save_debug_images,
generate_heatmaps=not self.generate_landmarks,
generate_landmarks=self.generate_landmarks,
translate_by_random_factor=False)
self.dataset_train = dataset.dataset_train()
self.dataset_val = dataset.dataset_val()
networks = {'scn': network_scn, 'unet': network_unet}
self.network = networks[self.network_id]
self.point_statistics_names = [
'pe_mean', 'pe_stdev', 'pe_median', 'num_correct'
]
self.additional_summaries_placeholders_val = dict([
(name, create_summary_placeholder(name))
for name in self.point_statistics_names
])
def loss_function(self, pred, target):
batch_size, _, _ = get_batch_channel_image_size(pred, self.data_format)
return tf.nn.l2_loss(pred - target) / batch_size
def initNetworks(self):
net = tf.make_template('scn', self.network)
if self.data_format == 'channels_first':
if self.generate_landmarks:
data_generator_entries = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks, 4])
])
else:
data_generator_entries = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('heatmaps',
[self.num_landmarks] + list(reversed(self.heatmap_size)))
])
else:
if self.generate_landmarks:
data_generator_entries = OrderedDict([
('image',
list(reversed(self.image_size)) + [self.image_channels]),
('landmarks', [self.num_landmarks, 4])
])
else:
data_generator_entries = OrderedDict([
('image',
list(reversed(self.image_size)) + [self.image_channels]),
('heatmaps',
list(reversed(self.heatmap_size)) + [self.num_landmarks])
])
sigmas = tf.get_variable('sigmas', [self.num_landmarks],
initializer=tf.constant_initializer(
self.heatmap_sigma))
mean_sigmas = tf.reduce_mean(sigmas)
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size)
#self.train_queue = DataGeneratorDataset(self.dataset_train, data_generator_entries, batch_size=self.batch_size)
placeholders = self.train_queue.dequeue()
image = placeholders[0]
if self.generate_landmarks:
target_landmarks = placeholders[1]
target_heatmaps = generate_heatmap_target(
list(reversed(self.heatmap_size)),
target_landmarks,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
loss_sigmas = self.sigma_regularization * tf.nn.l2_loss(
sigmas * target_landmarks[0, :, 0])
else:
target_heatmaps = placeholders[1]
loss_sigmas = self.sigma_regularization * tf.nn.l2_loss(sigmas)
heatmaps, _, _ = net(image,
num_heatmaps=self.num_landmarks,
is_training=True,
data_format=self.data_format)
self.loss_net = self.loss_function(heatmaps, target_heatmaps)
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
if self.reg_constant > 0:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg + loss_sigmas
else:
self.loss = self.loss_net + loss_sigmas
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg),
('loss_sigmas', loss_sigmas),
('mean_sigmas', mean_sigmas)])
# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss)
self.optimizer = tf.train.MomentumOptimizer(
learning_rate=self.learning_rate, momentum=0.99,
use_nesterov=True).minimize(self.loss)
# build val graph
val_placeholders = create_placeholders(data_generator_entries,
shape_prefix=[1])
self.image_val = val_placeholders['image']
if self.generate_landmarks:
self.target_landmarks_val = val_placeholders['landmarks']
self.target_heatmaps_val = generate_heatmap_target(
list(reversed(self.heatmap_size)),
self.target_landmarks_val,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
loss_sigmas_val = self.sigma_regularization * tf.nn.l2_loss(
sigmas * self.target_landmarks_val[0, :, 0])
else:
self.target_heatmaps_val = val_placeholders['heatmaps']
loss_sigmas_val = self.sigma_regularization * tf.nn.l2_loss(sigmas)
self.heatmaps_val, self.heatmals_local_val, self.heatmaps_global_val = net(
self.image_val,
num_heatmaps=self.num_landmarks,
is_training=False,
data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.heatmaps_val,
self.target_heatmaps_val)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg),
('loss_sigmas', loss_sigmas_val),
('mean_sigmas', mean_sigmas)])
def test_full_image(self, dataset_entry):
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
heatmap_transform = transformations['image']
feed_dict = {
self.image_val: np.expand_dims(generators['image'], axis=0)
}
if self.generate_landmarks:
feed_dict[self.target_landmarks_val] = np.expand_dims(
generators['landmarks'], axis=0)
else:
feed_dict[self.target_heatmaps_val] = np.expand_dims(
generators['heatmaps'], axis=0)
# run loss and update loss accumulators
run_tuple = self.sess.run(
(self.heatmaps_val, self.target_heatmaps_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(), feed_dict)
heatmaps = np.squeeze(run_tuple[0], axis=0)
image = generators['image']
return image, heatmaps, heatmap_transform
def test_cropped_image(self, dataset_entry):
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
heatmap_transform = transformations['image']
image_size_np = [1] + list(reversed(self.image_size))
heatmap_size_np = [self.num_landmarks] + list(reversed(
self.image_size))
full_image = generators['image']
landmarks = generators['landmarks']
image_tiler = ImageTiler(full_image.shape, image_size_np,
self.cropped_inc, True, -1)
landmark_tiler = LandmarkTiler(full_image.shape, image_size_np,
self.cropped_inc)
heatmap_tiler = ImageTiler(
(self.num_landmarks, ) + full_image.shape[1:], heatmap_size_np,
self.cropped_inc, True, 0)
for image_tiler, landmark_tiler, heatmap_tiler in zip(
image_tiler, landmark_tiler, heatmap_tiler):
current_image = image_tiler.get_current_data(full_image)
current_landmarks = landmark_tiler.get_current_data(landmarks)
feed_dict = {
self.image_val:
np.expand_dims(current_image, axis=0),
self.target_landmarks_val:
np.expand_dims(current_landmarks, axis=0)
}
run_tuple = self.sess.run(
(self.heatmaps_val, self.target_heatmaps_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(), feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
image_tiler.set_current_data(current_image)
heatmap_tiler.set_current_data(prediction)
return image_tiler.output_image, heatmap_tiler.output_image, heatmap_transform
def test(self):
print('Testing...')
if self.data_format == 'channels_first':
np_channel_index = 0
else:
np_channel_index = 3
heatmap_maxima = HeatmapTest(np_channel_index, False)
landmark_statistics = LandmarkStatistics()
landmarks = {}
for i in range(self.dataset_val.num_entries()):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
image_datasource = datasources['image_datasource']
landmarks_datasource = datasources['landmarks_datasource']
if not self.cropped_training:
image, heatmaps, heatmap_transform = self.test_full_image(
dataset_entry)
else:
image, heatmaps, heatmap_transform = self.test_cropped_image(
dataset_entry)
utils.io.image.write_np(
ShiftScaleClamp(scale=255, clamp_min=0,
clamp_max=255)(heatmaps).astype(np.uint8),
self.output_file_for_current_iteration(current_id +
'_heatmaps.mha'))
utils.io.image.write_np(
image,
self.output_file_for_current_iteration(current_id +
'_image.mha'))
predicted_landmarks = heatmap_maxima.get_landmarks(
heatmaps, image_datasource, self.image_spacing,
heatmap_transform)
landmarks[current_id] = predicted_landmarks
landmark_statistics.add_landmarks(current_id, predicted_landmarks,
landmarks_datasource)
tensorflow_train.utils.tensorflow_util.print_progress_bar(
i,
self.dataset_val.num_entries(),
prefix='Testing ',
suffix=' complete')
tensorflow_train.utils.tensorflow_util.print_progress_bar(
self.dataset_val.num_entries(),
self.dataset_val.num_entries(),
prefix='Testing ',
suffix=' complete')
print(landmark_statistics.get_pe_overview_string())
print(landmark_statistics.get_correct_id_string(20.0))
summary_values = OrderedDict(
zip(
self.point_statistics_names,
list(landmark_statistics.get_pe_statistics()) +
[landmark_statistics.get_correct_id(20)]))
# finalize loss values
self.val_loss_aggregator.finalize(self.current_iter, summary_values)
utils.io.landmark.save_points_csv(
landmarks, self.output_file_for_current_iteration('points.csv'))
overview_string = landmark_statistics.get_overview_string(
[2, 2.5, 3, 4, 10, 20], 10, 20.0)
utils.io.text.save_string_txt(
overview_string,
self.output_file_for_current_iteration('eval.txt'))
def initNetworks(self):
net = tf.make_template('scn', self.network)
if self.data_format == 'channels_first':
if self.generate_landmarks:
data_generator_entries = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks, 4])
])
else:
data_generator_entries = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('heatmaps',
[self.num_landmarks] + list(reversed(self.heatmap_size)))
])
else:
if self.generate_landmarks:
data_generator_entries = OrderedDict([
('image',
list(reversed(self.image_size)) + [self.image_channels]),
('landmarks', [self.num_landmarks, 4])
])
else:
data_generator_entries = OrderedDict([
('image',
list(reversed(self.image_size)) + [self.image_channels]),
('heatmaps',
list(reversed(self.heatmap_size)) + [self.num_landmarks])
])
sigmas = tf.get_variable('sigmas', [self.num_landmarks],
initializer=tf.constant_initializer(
self.heatmap_sigma))
mean_sigmas = tf.reduce_mean(sigmas)
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size)
#self.train_queue = DataGeneratorDataset(self.dataset_train, data_generator_entries, batch_size=self.batch_size)
placeholders = self.train_queue.dequeue()
image = placeholders[0]
if self.generate_landmarks:
target_landmarks = placeholders[1]
target_heatmaps = generate_heatmap_target(
list(reversed(self.heatmap_size)),
target_landmarks,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
loss_sigmas = self.sigma_regularization * tf.nn.l2_loss(
sigmas * target_landmarks[0, :, 0])
else:
target_heatmaps = placeholders[1]
loss_sigmas = self.sigma_regularization * tf.nn.l2_loss(sigmas)
heatmaps, _, _ = net(image,
num_heatmaps=self.num_landmarks,
is_training=True,
data_format=self.data_format)
self.loss_net = self.loss_function(heatmaps, target_heatmaps)
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
if self.reg_constant > 0:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg + loss_sigmas
else:
self.loss = self.loss_net + loss_sigmas
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg),
('loss_sigmas', loss_sigmas),
('mean_sigmas', mean_sigmas)])
# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss)
self.optimizer = tf.train.MomentumOptimizer(
learning_rate=self.learning_rate, momentum=0.99,
use_nesterov=True).minimize(self.loss)
# build val graph
val_placeholders = create_placeholders(data_generator_entries,
shape_prefix=[1])
self.image_val = val_placeholders['image']
if self.generate_landmarks:
self.target_landmarks_val = val_placeholders['landmarks']
self.target_heatmaps_val = generate_heatmap_target(
list(reversed(self.heatmap_size)),
self.target_landmarks_val,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
loss_sigmas_val = self.sigma_regularization * tf.nn.l2_loss(
sigmas * self.target_landmarks_val[0, :, 0])
else:
self.target_heatmaps_val = val_placeholders['heatmaps']
loss_sigmas_val = self.sigma_regularization * tf.nn.l2_loss(sigmas)
self.heatmaps_val, self.heatmals_local_val, self.heatmaps_global_val = net(
self.image_val,
num_heatmaps=self.num_landmarks,
is_training=False,
data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.heatmaps_val,
self.target_heatmaps_val)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg),
('loss_sigmas', loss_sigmas_val),
('mean_sigmas', mean_sigmas)])
class MainLoop(MainLoopBase):
def __init__(self, network_id, cv, landmark_source, sigma_regularization, output_folder_name=''):
super().__init__()
self.network_id = network_id
self.output_folder = os.path.join('output', network_id, landmark_source, cv if cv >= 0 else 'all', output_folder_name, self.output_folder_timestamp())
self.batch_size = 1
self.max_iter = 30000
self.learning_rate = 0.000001
self.test_iter = 5000
self.disp_iter = 100
self.snapshot_iter = self.test_iter
self.test_initialization = False
self.current_iter = 0
self.reg_constant = 0.001
self.cv = cv
self.landmark_source = landmark_source
original_image_extend = [193.5, 240.0]
image_sizes = {'unet': [512, 512],
'scn_mia': [512, 512]}
heatmap_sizes = {'unet': [512, 512],
'scn_mia': [512, 512]}
sigmas = {'unet': 2.5,
'scn_mia': 2.5}
self.image_size = image_sizes[self.network_id]
self.heatmap_size = heatmap_sizes[self.network_id]
self.image_spacing = [float(np.max([e / s for e, s in zip(original_image_extend, self.image_size)]))] * 2
self.sigma = sigmas[self.network_id]
self.image_channels = 1
self.num_landmarks = 19
self.heatmap_sigma = self.sigma
self.sigma_regularization = sigma_regularization
self.sigma_scale = 100.0
self.data_format = 'channels_first'
self.save_debug_images = False
self.base_folder = './'
dataset_parameters = {'image_size': self.image_size,
'heatmap_size': self.heatmap_size,
'image_spacing': self.image_spacing,
'num_landmarks': self.num_landmarks,
'base_folder': self.base_folder,
'data_format': self.data_format,
'save_debug_images': self.save_debug_images,
'cv': self.cv,
'landmark_source': self.landmark_source}
dataset = Dataset(**dataset_parameters)
self.dataset_train = dataset.dataset_train()
self.dataset_val = dataset.dataset_val()
networks = {'unet': network_unet,
'scn_mia': network_scn_mia}
self.network = networks[self.network_id]
self.landmark_metrics = ['pe_mean', 'pe_std', 'pe_median', 'or2', 'or25', 'or3', 'or4', 'or10']
self.landmark_metric_prefixes = ['challenge', 'senior', 'junior', 'mean']
self.additional_summaries_placeholders_val = OrderedDict([(prefix + '_' + name, create_summary_placeholder(prefix + '_' + name)) for name in self.landmark_metrics for prefix in self.landmark_metric_prefixes])
def loss_function(self, target, prediction):
return tf.nn.l2_loss(target - prediction) / get_batch_channel_image_size(target, self.data_format)[0]
def loss_sigmas(self, sigmas, landmarks):
return self.sigma_regularization * tf.nn.l2_loss(sigmas[None, :] * landmarks[:, :, 0]) / landmarks.get_shape().as_list()[0]
def initNetworks(self):
net = tf.make_template('net', self.network)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([('image', [self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks, 3])])
data_generator_entries_val = OrderedDict([('image', [self.image_channels] + list(reversed(self.image_size))),
('landmarks', [self.num_landmarks, 3])])
else:
data_generator_entries = OrderedDict([('image', list(reversed(self.image_size)) + [self.image_channels]),
('landmarks', [self.num_landmarks, 3])])
data_generator_entries_val = OrderedDict([('image', list(reversed(self.image_size)) + [self.image_channels]),
('landmarks', [self.num_landmarks, 3])])
sigmas = tf.get_variable('sigmas', [self.num_landmarks], initializer=tf.constant_initializer(self.heatmap_sigma))
sigmas_list = [(f's{i}', sigmas[i]) for i in range(self.num_landmarks)]
# build training graph
self.train_queue = DataGenerator(self.dataset_train, self.coord, data_generator_entries, batch_size=self.batch_size, n_threads=8)
placeholders = self.train_queue.dequeue()
image = placeholders[0]
target_landmarks = placeholders[1]
prediction = net(image, num_landmarks=self.num_landmarks, is_training=True, data_format=self.data_format)
target_heatmaps = generate_heatmap_target(list(reversed(self.heatmap_size)), target_landmarks, sigmas, scale=self.sigma_scale, normalize=True, data_format=self.data_format)
loss_sigmas = self.loss_sigmas(sigmas, target_landmarks)
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss_net = self.loss_function(target_heatmaps, prediction)
self.loss = self.loss_net + tf.cast(self.loss_reg, tf.float32) + loss_sigmas
# optimizer
global_step = tf.Variable(self.current_iter, trainable=False)
optimizer = tf.train.MomentumOptimizer(learning_rate=self.learning_rate, momentum=0.99, use_nesterov=True)
unclipped_gradients, variables = zip(*optimizer.compute_gradients(self.loss))
norm = tf.global_norm(unclipped_gradients)
gradients, _ = tf.clip_by_global_norm(unclipped_gradients, 10000.0)
self.optimizer = optimizer.apply_gradients(zip(gradients, variables), global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net), ('loss_reg', self.loss_reg), ('loss_sigmas', loss_sigmas), ('norm', norm)] + sigmas_list)
# build val graph
self.val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(data_generator_entries_val, shape_prefix=[1])
self.image_val = self.val_placeholders['image']
self.target_landmarks_val = self.val_placeholders['landmarks']
self.prediction_val = net(self.image_val, num_landmarks=self.num_landmarks, is_training=False, data_format=self.data_format)
self.target_heatmaps_val = generate_heatmap_target(list(reversed(self.heatmap_size)), self.target_landmarks_val, sigmas, scale=self.sigma_scale, normalize=True, data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.target_heatmaps_val, self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val), ('loss_reg', self.loss_reg), ('loss_sigmas', tf.constant(0, tf.float32)), ('norm', tf.constant(0, tf.float32))] + sigmas_list)
def test_full_image(self, dataset_entry):
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
feed_dict = {self.val_placeholders['image']: np.expand_dims(generators['image'], axis=0),
self.val_placeholders['landmarks']: np.expand_dims(generators['landmarks'], axis=0)}
# run loss and update loss accumulators
run_tuple = self.sess.run((self.prediction_val, self.target_heatmaps_val, self.loss_val) + self.val_loss_aggregator.get_update_ops(), feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
target_heatmaps = np.squeeze(run_tuple[1], axis=0)
image = generators['image']
transformation = transformations['image']
return image, prediction, target_heatmaps, transformation
def finalize_landmark_statistics(self, landmark_statistics, prefix):
pe_mean, pe_std, pe_median = landmark_statistics.get_pe_statistics()
or2, or25, or3, or4, or10 = landmark_statistics.get_num_outliers([2.0, 2.5, 3.0, 4.0, 10.0], True)
print(prefix + '_pe', ['{0:.3f}'.format(s) for s in [pe_mean, pe_std, pe_median]])
print(prefix + '_outliers', ['{0:.3f}'.format(s) for s in [or2, or25, or3, or4, or10]])
overview_string = landmark_statistics.get_overview_string([2, 2.5, 3, 4, 10, 20], 10, 20.0)
utils.io.text.save_string_txt(overview_string, self.output_file_for_current_iteration(prefix + '_eval.txt'))
additional_summaries = {prefix + '_pe_mean': pe_mean,
prefix + '_pe_std': pe_std,
prefix + '_pe_median': pe_median,
prefix + '_or2': or2,
prefix + '_or25': or25,
prefix + '_or3': or3,
prefix + '_or4': or4,
prefix + '_or10': or10}
return additional_summaries
def test(self):
heatmap_test = HeatmapTest(channel_axis=0, invert_transformation=False)
challenge_landmark_statistics = LandmarkStatistics()
senior_landmark_statistics = LandmarkStatistics()
junior_landmark_statistics = LandmarkStatistics()
mean_landmark_statistics = LandmarkStatistics()
landmarks = {}
for i in range(self.dataset_val.num_entries()):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
reference_image = datasources['image_datasource']
groundtruth_challenge_landmarks = datasources['challenge_landmarks_datasource']
groundtruth_senior_landmarks = datasources['senior_landmarks_datasource']
groundtruth_junior_landmarks = datasources['junior_landmarks_datasource']
groundtruth_mean_landmarks = datasources['mean_landmarks_datasource']
image, prediction, target_heatmaps, transform = self.test_full_image(dataset_entry)
utils.io.image.write_np(image, self.output_file_for_current_iteration(current_id + '_image.mha'))
utils.io.image.write_np(prediction, self.output_file_for_current_iteration(current_id + '_prediction.mha'))
utils.io.image.write_np(target_heatmaps, self.output_file_for_current_iteration(current_id + '_target_heatmap.mha'))
predicted_landmarks = heatmap_test.get_landmarks(prediction, reference_image, output_spacing=self.image_spacing, transformation=transform)
tensorflow_train.utils.tensorflow_util.print_progress_bar(i, self.dataset_val.num_entries())
landmarks[current_id] = predicted_landmarks
challenge_landmark_statistics.add_landmarks(current_id, predicted_landmarks, groundtruth_challenge_landmarks)
senior_landmark_statistics.add_landmarks(current_id, predicted_landmarks, groundtruth_senior_landmarks)
junior_landmark_statistics.add_landmarks(current_id, predicted_landmarks, groundtruth_junior_landmarks)
mean_landmark_statistics.add_landmarks(current_id, predicted_landmarks, groundtruth_mean_landmarks)
tensorflow_train.utils.tensorflow_util.print_progress_bar(self.dataset_val.num_entries(), self.dataset_val.num_entries())
challenge_summaries = self.finalize_landmark_statistics(challenge_landmark_statistics, 'challenge')
senior_summaries = self.finalize_landmark_statistics(senior_landmark_statistics, 'senior')
junior_summaries = self.finalize_landmark_statistics(junior_landmark_statistics, 'junior')
mean_summaries = self.finalize_landmark_statistics(mean_landmark_statistics, 'mean')
additional_summaries = OrderedDict(chain(senior_summaries.items(), junior_summaries.items(), challenge_summaries.items(), mean_summaries.items()))
# finalize loss values
self.val_loss_aggregator.finalize(self.current_iter, additional_summaries)
utils.io.landmark.save_points_csv(landmarks, self.output_file_for_current_iteration('prediction.csv'))
class MainLoop(MainLoopBase):
def __init__(self,
cv,
network,
unet,
network_parameters,
learning_rate,
output_folder_name=''):
"""
Initializer.
:param cv: The cv fold. 0, 1, 2 for CV; 'train_all' for training on whole dataset.
:param network: The used network. Usually network_u.
:param unet: The specific instance of the U-Net. Usually UnetClassicAvgLinear3d.
:param network_parameters: The network parameters passed to unet.
:param learning_rate: The initial learning rate.
:param output_folder_name: The output folder name that is used for distinguishing experiments.
"""
super().__init__()
self.batch_size = 1
self.learning_rates = [
learning_rate, learning_rate * 0.5, learning_rate * 0.1
]
self.learning_rate_boundaries = [10000, 15000]
self.max_iter = 20000
self.test_iter = 5000
self.disp_iter = 100
self.snapshot_iter = 5000
self.test_initialization = False
self.current_iter = 0
self.reg_constant = 0.0005
self.use_background = True
self.num_labels = 1
self.heatmap_sigma = 2.0
self.data_format = 'channels_first'
self.network = network
self.unet = unet
self.network_parameters = network_parameters
self.padding = 'same'
self.use_pyro_dataset = False
self.save_output_images = True
self.save_output_images_as_uint = True # set to False, if you want to see the direct network output
self.save_debug_images = False
self.has_validation_groundtruth = cv in [0, 1, 2]
self.local_base_folder = '../verse2019_dataset'
self.image_size = [64, 64, 128]
self.image_spacing = [8] * 3
self.output_folder = os.path.join('./output/spine_localization/',
network.__name__, unet.__name__,
output_folder_name, str(cv),
self.output_folder_timestamp())
dataset_parameters = {
'base_folder': self.local_base_folder,
'image_size': self.image_size,
'image_spacing': self.image_spacing,
'cv': cv,
'input_gaussian_sigma': 3.0,
'generate_spine_heatmap': True,
'save_debug_images': self.save_debug_images
}
dataset = Dataset(**dataset_parameters)
if self.use_pyro_dataset:
server_name = '@localhost:51232'
uri = 'PYRO:verse_dataset' + server_name
print('using pyro uri', uri)
self.dataset_train = PyroClientDataset(uri, **dataset_parameters)
else:
self.dataset_train = dataset.dataset_train()
self.dataset_val = dataset.dataset_val()
self.point_statistics_names = ['pe_mean', 'pe_stdev', 'pe_median']
self.additional_summaries_placeholders_val = dict([
(name, create_summary_placeholder(name))
for name in self.point_statistics_names
])
def loss_function(self, pred, target):
"""
L2 loss function calculated with prediction and target.
:param pred: The predicted image.
:param target: The target image.
:return: L2 loss of (pred - target) / batch_size
"""
batch_size, _, _ = get_batch_channel_image_size(pred, self.data_format)
return tf.nn.l2_loss(pred - target) / batch_size
def init_networks(self):
"""
Initialize networks and placeholders.
"""
network_image_size = list(reversed(self.image_size))
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image', [1] + network_image_size),
('spine_heatmap', [1] + network_image_size)
])
else:
data_generator_entries = OrderedDict([
('image', network_image_size + [1]),
('spine_heatmap', [1] + network_image_size)
])
data_generator_types = {
'image': tf.float32,
'spine_heatmap': tf.float32
}
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(
coord=self.coord,
dataset=self.dataset_train,
data_names_and_shapes=data_generator_entries,
data_types=data_generator_types,
batch_size=self.batch_size)
data, target_spine_heatmap = self.train_queue.dequeue()
prediction = training_net(data,
num_labels=self.num_labels,
is_training=True,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
self.loss_net = self.loss_function(target=target_spine_heatmap,
pred=prediction)
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss = self.loss_net + tf.cast(self.loss_reg, tf.float32)
# solver
global_step = tf.Variable(self.current_iter, trainable=False)
learning_rate = tf.train.piecewise_constant(
global_step, self.learning_rate_boundaries, self.learning_rates)
self.optimizer = tf.train.AdamOptimizer(
learning_rate=learning_rate).minimize(self.loss,
global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg)])
# build val graph
self.data_val, self.target_spine_heatmap_val = create_placeholders_tuple(
data_generator_entries,
data_types=data_generator_types,
shape_prefix=[1])
self.prediction_val = training_net(self.data_val,
num_labels=self.num_labels,
is_training=False,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
if self.has_validation_groundtruth:
self.loss_val = self.loss_function(
target=self.target_spine_heatmap_val, pred=self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg)])
def test_full_image(self, dataset_entry):
"""
Perform inference on a dataset_entry with the validation network.
:param dataset_entry: A dataset entry from the dataset.
:return: input image (np.array), network prediction (np.array), transformation (sitk.Transform)
"""
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
if self.has_validation_groundtruth:
feed_dict = {
self.data_val:
np.expand_dims(generators['image'], axis=0),
self.target_spine_heatmap_val:
np.expand_dims(generators['spine_heatmap'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run(
(self.prediction_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(),
feed_dict=feed_dict)
else:
feed_dict = {
self.data_val: np.expand_dims(generators['image'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run((self.prediction_val, ),
feed_dict=feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
transformation = transformations['image']
image = generators['image']
return image, prediction, transformation
def test(self):
"""
The test function. Performs inference on the the validation images and calculates the loss.
"""
print('Testing...')
channel_axis = 0
if self.data_format == 'channels_last':
channel_axis = 3
landmark_statistics = LandmarkStatistics()
landmarks = {}
num_entries = self.dataset_val.num_entries()
for i in range(num_entries):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
if self.has_validation_groundtruth:
groundtruth_landmarks = datasources['landmarks']
groundtruth_landmark = [
get_mean_landmark(groundtruth_landmarks)
]
input_image = datasources['image']
image, prediction, transformation = self.test_full_image(
dataset_entry)
predictions_sitk = utils.sitk_image.transform_np_output_to_sitk_input(
output_image=prediction,
output_spacing=self.image_spacing,
channel_axis=channel_axis,
input_image_sitk=input_image,
transform=transformation,
interpolator='linear',
output_pixel_type=sitk.sitkFloat32)
if self.save_output_images:
if self.save_output_images_as_uint:
image_normalization = 'min_max'
heatmap_normalization = (0, 1)
output_image_type = np.uint8
else:
image_normalization = None
heatmap_normalization = None
output_image_type = np.float32
origin = transformation.TransformPoint(np.zeros(3, np.float64))
utils.io.image.write_multichannel_np(
image,
self.output_file_for_current_iteration(current_id +
'_input.mha'),
output_normalization_mode=image_normalization,
data_format=self.data_format,
image_type=output_image_type,
spacing=self.image_spacing,
origin=origin)
utils.io.image.write_multichannel_np(
prediction,
self.output_file_for_current_iteration(current_id +
'_prediction.mha'),
output_normalization_mode=heatmap_normalization,
data_format=self.data_format,
image_type=output_image_type,
spacing=self.image_spacing,
origin=origin)
#utils.io.image.write(predictions_sitk[0], self.output_file_for_current_iteration(current_id + '_prediction_original.mha'))
predictions_com = input_image.TransformContinuousIndexToPhysicalPoint(
list(
reversed(
utils.np_image.center_of_mass(
utils.sitk_np.sitk_to_np_no_copy(
predictions_sitk[0])))))
current_landmark = [Landmark(predictions_com)]
landmarks[current_id] = current_landmark
if self.has_validation_groundtruth:
landmark_statistics.add_landmarks(current_id, current_landmark,
groundtruth_landmark)
print_progress_bar(i,
num_entries,
prefix='Testing ',
suffix=' complete')
utils.io.landmark.save_points_csv(
landmarks, self.output_file_for_current_iteration('points.csv'))
# finalize loss values
if self.has_validation_groundtruth:
print(landmark_statistics.get_pe_overview_string())
summary_values = OrderedDict(
zip(self.point_statistics_names,
list(landmark_statistics.get_pe_statistics())))
# finalize loss values
self.val_loss_aggregator.finalize(self.current_iter,
summary_values)
overview_string = landmark_statistics.get_overview_string(
[2, 2.5, 3, 4, 10, 20], 10)
utils.io.text.save_string_txt(
overview_string,
self.output_file_for_current_iteration('eval.txt'))
def init_networks(self):
network_image_size = self.image_size
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image', [1] + network_image_size),
('labels', [self.num_labels] + network_image_size)
])
else:
data_generator_entries = OrderedDict([
('image', network_image_size + [1]),
('labels', network_image_size + [self.num_labels])
])
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size)
data, mask = self.train_queue.dequeue()
prediction, _, _ = training_net(data,
num_labels=self.num_labels,
is_training=True,
data_format=self.data_format)
# losses
self.loss_net = self.loss_function(labels=mask,
logits=prediction,
data_format=self.data_format)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
if self.reg_constant > 0:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg
else:
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg)])
# solver
global_step = tf.Variable(self.current_iter)
learning_rate = tf.train.piecewise_constant(
global_step, self.learning_rate_boundaries, self.learning_rates)
self.optimizer = tf.train.AdamOptimizer(
learning_rate=learning_rate).minimize(self.loss,
global_step=global_step)
# build val graph
val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(
data_generator_entries, shape_prefix=[1])
self.data_val = val_placeholders['image']
self.prediction_val, self.local_prediction_val, self.spatial_prediction_val = training_net(
self.data_val,
num_labels=self.num_labels,
is_training=False,
data_format=self.data_format)
self.prediction_softmax_val = tf.nn.softmax(
self.prediction_val,
axis=1 if self.data_format == 'channels_first' else 4)
if self.has_validation_groundtruth:
self.mask_val = val_placeholders['labels']
# losses
self.loss_val = self.loss_function(labels=self.mask_val,
logits=self.prediction_val,
data_format=self.data_format)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg)])
class MainLoop(MainLoopBase):
def __init__(self, modality, cv):
super().__init__()
self.modality = modality
self.cv = cv
self.batch_size = 1
self.learning_rates = [0.00001, 0.000001]
self.learning_rate_boundaries = [20000]
self.max_iter = 40000
self.test_iter = 5000
self.disp_iter = 100
self.snapshot_iter = self.test_iter
self.test_initialization = False
self.current_iter = 0
self.reg_constant = 0.0001
self.num_labels = 8
self.data_format = 'channels_first'
self.channel_axis = 1
self.save_debug_images = False
self.has_validation_groundtruth = cv != 0
self.base_folder = 'mmwhs_dataset'
self.image_size = [64, 64, 64]
if modality == 'ct':
self.image_spacing = [3, 3, 3]
else:
self.image_spacing = [4, 4, 4]
self.input_gaussian_sigma = 1.0
self.label_gaussian_sigma = 1.0
self.use_landmarks = True
self.output_folder = './output/scn_' + modality + '_' + str(
cv) + '/' + self.output_folder_timestamp()
dataset_parameters = dict(
image_size=self.image_size,
image_spacing=self.image_spacing,
base_folder=self.base_folder,
cv=self.cv,
modality=self.modality,
input_gaussian_sigma=self.input_gaussian_sigma,
label_gaussian_sigma=self.label_gaussian_sigma,
use_landmarks=self.use_landmarks,
num_labels=self.num_labels,
data_format=self.data_format,
save_debug_images=self.save_debug_images)
self.dataset = Dataset(**dataset_parameters)
self.dataset_train = self.dataset.dataset_train()
self.dataset_val = self.dataset.dataset_val()
self.dataset_val = self.dataset.dataset_val()
self.files_to_copy = ['main.py', 'network.py', 'dataset.py']
self.dice_names = list(
map(lambda x: 'dice_{}'.format(x), range(self.num_labels)))
self.additional_summaries_placeholders_val = dict([
(name, create_summary_placeholder(name))
for name in self.dice_names
])
self.loss_function = softmax_cross_entropy_with_logits
self.network = network_scn
def init_networks(self):
network_image_size = self.image_size
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image', [1] + network_image_size),
('labels', [self.num_labels] + network_image_size)
])
else:
data_generator_entries = OrderedDict([
('image', network_image_size + [1]),
('labels', network_image_size + [self.num_labels])
])
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size)
data, mask = self.train_queue.dequeue()
prediction, _, _ = training_net(data,
num_labels=self.num_labels,
is_training=True,
data_format=self.data_format)
# losses
self.loss_net = self.loss_function(labels=mask,
logits=prediction,
data_format=self.data_format)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
if self.reg_constant > 0:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg
else:
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg)])
# solver
global_step = tf.Variable(self.current_iter)
learning_rate = tf.train.piecewise_constant(
global_step, self.learning_rate_boundaries, self.learning_rates)
self.optimizer = tf.train.AdamOptimizer(
learning_rate=learning_rate).minimize(self.loss,
global_step=global_step)
# build val graph
val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(
data_generator_entries, shape_prefix=[1])
self.data_val = val_placeholders['image']
self.prediction_val, self.local_prediction_val, self.spatial_prediction_val = training_net(
self.data_val,
num_labels=self.num_labels,
is_training=False,
data_format=self.data_format)
self.prediction_softmax_val = tf.nn.softmax(
self.prediction_val,
axis=1 if self.data_format == 'channels_first' else 4)
if self.has_validation_groundtruth:
self.mask_val = val_placeholders['labels']
# losses
self.loss_val = self.loss_function(labels=self.mask_val,
logits=self.prediction_val,
data_format=self.data_format)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg)])
def test(self):
print('Testing...')
channel_axis = 0
if self.data_format == 'channels_last':
channel_axis = 3
labels = list(range(self.num_labels))
segmentation_test = SegmentationTest(labels,
channel_axis=channel_axis,
interpolator='cubic',
largest_connected_component=False,
all_labels_are_connected=False)
segmentation_statistics = SegmentationStatistics(
labels,
self.output_folder_for_current_iteration(),
metrics={'dice': DiceMetric()})
num_entries = self.dataset_val.num_entries()
for i in range(num_entries):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
if self.has_validation_groundtruth:
feed_dict = {
self.data_val: np.expand_dims(generators['image'], axis=0),
self.mask_val: np.expand_dims(generators['labels'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run(
(self.prediction_softmax_val, self.local_prediction_val,
self.spatial_prediction_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(),
feed_dict=feed_dict)
else:
feed_dict = {
self.data_val: np.expand_dims(generators['data'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run((self.prediction_softmax_val, ),
feed_dict=feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
#local_prediction = np.squeeze(run_tuple[1], axis=0)
#spatial_prediction = np.squeeze(run_tuple[2], axis=0)
input = datasources['image']
transformation = transformations['image']
prediction_labels, prediction_sitk = segmentation_test.get_label_image(
prediction,
input,
self.image_spacing,
transformation,
return_transformed_sitk=True)
utils.io.image.write(
prediction_labels,
self.output_file_for_current_iteration(current_id + '.mha'))
origin = transformation.TransformPoint(np.zeros(3, np.float64))
utils.io.image.write_multichannel_np(
prediction,
self.output_file_for_current_iteration(current_id +
'_prediction.mha'),
output_normalization_mode=(0, 1),
data_format=self.data_format,
image_type=np.uint8,
spacing=self.image_spacing,
origin=origin)
#utils.io.image.write_multichannel_np(local_prediction, self.output_file_for_current_iteration(current_id + '_local_prediction.mha'), output_normalization_mode=(0, 1), data_format=self.data_format, image_type=np.uint8, spacing=self.image_spacing, origin=origin)
#utils.io.image.write_multichannel_np(spatial_prediction, self.output_file_for_current_iteration(current_id + '_spatial_prediction.mha'), output_normalization_mode=(0, 1), data_format=self.data_format, image_type=np.uint8, spacing=self.image_spacing, origin=origin)
if self.has_validation_groundtruth:
groundtruth = datasources['labels']
segmentation_statistics.add_labels(current_id,
prediction_labels,
groundtruth)
tensorflow_train.utils.tensorflow_util.print_progress_bar(
i, num_entries, prefix='Testing ', suffix=' complete')
# finalize loss values
if self.has_validation_groundtruth:
segmentation_statistics.finalize()
dice_list = segmentation_statistics.get_metric_mean_list('dice')
dice_dict = OrderedDict(list(zip(self.dice_names, dice_list)))
self.val_loss_aggregator.finalize(self.current_iter,
summary_values=dice_dict)
def initNetworks(self):
network_image_size = list(reversed(self.image_size))
global_step = tf.Variable(self.current_iter)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('data', [1] + network_image_size),
('mask', [self.num_labels] + network_image_size)
])
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size)
data, mask = self.train_queue.dequeue()
with tf.variable_scope('training', reuse=False):
if self.network.__name__ is 'network_ud' or self.network.__name__ is 'SegCaps_multilabels':
prediction = training_net(data,
num_labels=self.num_labels,
is_training=True,
data_format=self.data_format)
else:
prediction = training_net(data,
routing_type=self.routing_type,
num_labels=self.num_labels,
is_training=True,
data_format=self.data_format)
#print parameters count
logging.info('------------')
var_num = np.sum([
np.product([xi.value for xi in x.get_shape()])
for x in tf.global_variables()
])
logging.info('Net number of parameter : ' + str(var_num))
# losses
if 'spread_loss' in self.loss_function.__name__:
self.loss_net = self.loss_function(labels=mask,
logits=prediction,
global_step=global_step,
data_format=self.data_format)
else:
self.loss_net = self.loss_function(labels=mask,
logits=prediction,
data_format=self.data_format)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net)])
# solver
self.optimizer = tf.train.AdadeltaOptimizer(learning_rate=1).minimize(
self.loss,
global_step=global_step,
var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='training'))
# build val graph
val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(
data_generator_entries, shape_prefix=[1])
self.data_val = val_placeholders['data']
with tf.variable_scope('testing', reuse=True):
if self.network.__name__ is 'network_ud' or self.network.__name__ is 'SegCaps_multilabels':
self.prediction_val = training_net(
self.data_val,
num_labels=self.num_labels,
is_training=False,
data_format=self.data_format)
else:
self.prediction_val = training_net(
self.data_val,
routing_type=self.routing_type,
num_labels=self.num_labels,
is_training=False,
data_format=self.data_format)
self.mask_val = val_placeholders['mask']
# losses
self.loss_val = self.loss_function(labels=self.mask_val,
logits=self.prediction_val,
data_format=self.data_format)
self.val_losses = OrderedDict([('loss', self.loss_val)])
class MainLoop(MainLoopBase):
def __init__(self, param):
super().__init__()
#polyaxon
data_dir = os.path.join(
list(get_data_paths().values())[0], "lung/JSRT/preprocessed/")
logging.info('DATA DIR = ' + data_dir)
output_path = get_outputs_path()
self.loss_function = param[0]
self.network = param[1]
self.routing_type = param[2]
self.batch_size = 1
self.learning_rates = [1, 1]
self.max_iter = 300000
self.test_iter = 10000
self.disp_iter = 100
self.snapshot_iter = self.test_iter
self.test_initialization = False
self.current_iter = 0
self.num_labels = 6
self.data_format = 'channels_first' #WARNING: Capsule might not work with channel last !
self.channel_axis = 1
self.save_debug_images = False
self.base_folder = data_dir ##input folder
self.image_size = [128, 128]
self.image_spacing = [1, 1]
self.output_folder = output_path + self.network.__name__ + '_' + self.output_folder_timestamp(
) ##output save
self.dataset = Dataset(image_size=self.image_size,
image_spacing=self.image_spacing,
num_labels=self.num_labels,
base_folder=self.base_folder,
data_format=self.data_format,
save_debug_images=self.save_debug_images)
self.dataset_train = self.dataset.dataset_train()
self.dataset_train.get_next()
self.dataset_val = self.dataset.dataset_val()
self.dice_names = list(
map(lambda x: 'dice_{}'.format(x), range(self.num_labels)))
self.additional_summaries_placeholders_val = dict([
(name, create_summary_placeholder(name))
for name in self.dice_names
])
if self.network.__name__ is 'network_ud':
self.net_file = './Lung_Segmentation/LungSeg/cnn_network.py'
elif self.network.__name__ is 'SegCaps_multilabels':
self.net_file = './Lung_Segmentation/LungSeg/SegCaps/SegCaps.py'
else:
self.net_file = './Lung_Segmentation/LungSeg/capsule_network.py'
self.files_to_copy = ['main_train_and_test.py', self.net_file]
def initNetworks(self):
network_image_size = list(reversed(self.image_size))
global_step = tf.Variable(self.current_iter)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('data', [1] + network_image_size),
('mask', [self.num_labels] + network_image_size)
])
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size)
data, mask = self.train_queue.dequeue()
with tf.variable_scope('training', reuse=False):
if self.network.__name__ is 'network_ud' or self.network.__name__ is 'SegCaps_multilabels':
prediction = training_net(data,
num_labels=self.num_labels,
is_training=True,
data_format=self.data_format)
else:
prediction = training_net(data,
routing_type=self.routing_type,
num_labels=self.num_labels,
is_training=True,
data_format=self.data_format)
#print parameters count
logging.info('------------')
var_num = np.sum([
np.product([xi.value for xi in x.get_shape()])
for x in tf.global_variables()
])
logging.info('Net number of parameter : ' + str(var_num))
# losses
if 'spread_loss' in self.loss_function.__name__:
self.loss_net = self.loss_function(labels=mask,
logits=prediction,
global_step=global_step,
data_format=self.data_format)
else:
self.loss_net = self.loss_function(labels=mask,
logits=prediction,
data_format=self.data_format)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net)])
# solver
self.optimizer = tf.train.AdadeltaOptimizer(learning_rate=1).minimize(
self.loss,
global_step=global_step,
var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='training'))
# build val graph
val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(
data_generator_entries, shape_prefix=[1])
self.data_val = val_placeholders['data']
with tf.variable_scope('testing', reuse=True):
if self.network.__name__ is 'network_ud' or self.network.__name__ is 'SegCaps_multilabels':
self.prediction_val = training_net(
self.data_val,
num_labels=self.num_labels,
is_training=False,
data_format=self.data_format)
else:
self.prediction_val = training_net(
self.data_val,
routing_type=self.routing_type,
num_labels=self.num_labels,
is_training=False,
data_format=self.data_format)
self.mask_val = val_placeholders['mask']
# losses
self.loss_val = self.loss_function(labels=self.mask_val,
logits=self.prediction_val,
data_format=self.data_format)
self.val_losses = OrderedDict([('loss', self.loss_val)])
def test(self):
logging.info('Testing...')
channel_axis = 0
if self.data_format == 'channels_last':
channel_axis = 3
labels = list(range(self.num_labels))
segmentation_test = SegmentationTest(labels,
channel_axis=channel_axis,
interpolator='cubic',
largest_connected_component=False,
all_labels_are_connected=False)
segmentation_statistics = SegmentationStatistics(
labels,
self.output_folder_for_current_iteration(),
metrics={'dice': DiceMetric()})
num_entries = self.dataset_val.num_entries()
for i in range(num_entries):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
feed_dict = {
self.data_val: np.expand_dims(generators['data'], axis=0),
self.mask_val: np.expand_dims(generators['mask'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run(
(self.prediction_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(),
feed_dict=feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
input = datasources['image']
transformation = transformations['data']
prediction_labels, prediction_sitk = segmentation_test.get_label_image(
prediction,
input,
self.image_spacing,
transformation,
return_transformed_sitk=True)
utils.io.image.write_np_colormask(
prediction_labels,
self.output_file_for_current_iteration(current_id + '.png'))
utils.io.image.write_np(
prediction,
self.output_file_for_current_iteration(current_id +
'_prediction.mha'))
groundtruth = datasources['mask']
segmentation_statistics.add_labels(current_id, prediction_labels,
groundtruth)
tensorflow_train.utils.tensorflow_util.print_progress_bar(
i, num_entries, prefix='Testing ', suffix=' complete')
# finalize loss values
segmentation_statistics.finalize()
dice_list = segmentation_statistics.get_metric_mean_list('dice')
dice_dict = OrderedDict(list(zip(self.dice_names, dice_list)))
self.val_loss_aggregator.finalize(self.current_iter,
summary_values=dice_dict)
class MainLoop(MainLoopBase):
def __init__(self, cv, modality):
super().__init__()
self.cv = cv
self.output_folder = './mmwhs_localization/{}_{}'.format(
modality, cv) + '/' + self.output_folder_timestamp()
self.batch_size = 1
self.learning_rate = 0.00001
self.learning_rates = [self.learning_rate, self.learning_rate * 0.1]
self.learning_rate_boundaries = [10000]
self.max_iter = 20000
self.test_iter = 5000
self.disp_iter = 100
self.snapshot_iter = self.test_iter
self.test_initialization = False
self.current_iter = 0
self.reg_constant = 0.00005
self.invert_transformation = False
self.image_size = [32] * 3
if modality == 'ct':
self.image_spacing = [10] * 3
else:
self.image_spacing = [12] * 3
self.sigma = [1.5] * 3
self.image_channels = 1
self.num_landmarks = 1
self.data_format = 'channels_first'
self.save_debug_images = False
self.local_base_folder = '../../semantic_segmentation/mmwhs/mmwhs_dataset'
dataset_parameters = {
'base_folder': self.local_base_folder,
'image_size': self.image_size,
'image_spacing': self.image_spacing,
'cv': cv,
'input_gaussian_sigma': 4.0,
'modality': modality,
'save_debug_images': self.save_debug_images
}
dataset = Dataset(**dataset_parameters)
self.dataset_train = dataset.dataset_train()
self.dataset_val = dataset.dataset_val()
self.network = network_unet
self.loss_function = lambda x, y: tf.nn.l2_loss(
x - y) / get_batch_channel_image_size(x, self.data_format)[0]
def initNetworks(self):
net = tf.make_template('net', self.network)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks',
[self.num_landmarks] + list(reversed(self.image_size)))
])
data_generator_entries_val = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks',
[self.num_landmarks] + list(reversed(self.image_size)))
])
else:
raise NotImplementedError
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size)
image, landmarks = self.train_queue.dequeue()
prediction = net(image,
num_landmarks=self.num_landmarks,
is_training=True,
data_format=self.data_format)
self.loss_net = self.loss_function(landmarks, prediction)
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
if self.reg_constant > 0:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg
else:
self.loss_reg = 0
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg)])
global_step = tf.Variable(self.current_iter)
learning_rate = tf.train.piecewise_constant(
global_step, self.learning_rate_boundaries, self.learning_rates)
#self.optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=0.99, use_nesterov=True).minimize(self.loss, global_step=global_step)
self.optimizer = tf.train.AdamOptimizer(
learning_rate=learning_rate).minimize(self.loss,
global_step=global_step)
# build val graph
self.val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(
data_generator_entries_val, shape_prefix=[1])
self.image_val = self.val_placeholders['image']
self.landmarks_val = self.val_placeholders['landmarks']
self.prediction_val = net(self.image_val,
num_landmarks=self.num_landmarks,
is_training=False,
data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.landmarks_val,
self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg)])
def test_full_image(self, dataset_entry):
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
feed_dict = {
self.val_placeholders['image']:
np.expand_dims(generators['image'], axis=0),
self.val_placeholders['landmarks']:
np.expand_dims(generators['landmarks'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run((self.prediction_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(),
feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
image = generators['image']
transformation = transformations['image']
return image, prediction, transformation
def test(self):
heatmap_test = HeatmapTest(channel_axis=0, invert_transformation=False)
landmark_statistics = LandmarkStatistics()
landmarks = {}
for i in range(self.dataset_val.num_entries()):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
reference_image = datasources['image']
groundtruth_landmarks = datasources['landmarks']
image, prediction, transform = self.test_full_image(dataset_entry)
utils.io.image.write_np(
(prediction * 128).astype(np.int8),
self.output_file_for_current_iteration(current_id +
'_heatmap.mha'))
predicted_landmarks = heatmap_test.get_landmarks(
prediction,
reference_image,
transformation=transform,
output_spacing=self.image_spacing)
tensorflow_train.utils.tensorflow_util.print_progress_bar(
i, self.dataset_val.num_entries())
landmarks[current_id] = predicted_landmarks
landmark_statistics.add_landmarks(current_id, predicted_landmarks,
groundtruth_landmarks)
tensorflow_train.utils.tensorflow_util.print_progress_bar(
self.dataset_val.num_entries(), self.dataset_val.num_entries())
overview_string = landmark_statistics.get_overview_string(
[2.0, 4.0, 10.0])
print(overview_string)
# finalize loss values
self.val_loss_aggregator.finalize(self.current_iter)
utils.io.landmark.save_points_csv(
landmarks,
self.output_file_for_current_iteration('prediction.csv'))
utils.io.text.save_string_txt(
overview_string,
self.output_file_for_current_iteration('summary.txt'))
class MainLoop(MainLoopBase):
def __init__(self,
cv,
network,
unet,
network_parameters,
learning_rate,
output_folder_name=''):
"""
Initializer.
:param cv: The cv fold. 0, 1, 2 for CV; 'train_all' for training on whole dataset.
:param network: The used network. Usually network_u.
:param unet: The specific instance of the U-Net. Usually UnetClassicAvgLinear3d.
:param network_parameters: The network parameters passed to unet.
:param learning_rate: The initial learning rate.
:param output_folder_name: The output folder name that is used for distinguishing experiments.
"""
super().__init__()
self.batch_size = 1
self.learning_rates = [
learning_rate, learning_rate * 0.5, learning_rate * 0.1
]
self.learning_rate_boundaries = [50000, 75000]
self.max_iter = 100000
self.test_iter = 10000
self.disp_iter = 100
self.snapshot_iter = 5000
self.test_initialization = False
self.current_iter = 0
self.reg_constant = 0.0005
self.use_background = True
self.num_landmarks = 25
self.heatmap_sigma = 4.0
self.learnable_sigma = True
self.data_format = 'channels_first'
self.network = network
self.unet = unet
self.network_parameters = network_parameters
self.padding = 'same'
self.clip_gradient_global_norm = 100000.0
self.use_pyro_dataset = False
self.use_spine_postprocessing = True
self.save_output_images = True
self.save_output_images_as_uint = True # set to False, if you want to see the direct network output
self.save_debug_images = False
self.has_validation_groundtruth = cv in [0, 1, 2]
self.local_base_folder = '../verse2019_dataset'
self.image_size = [96, 96, 128]
self.image_spacing = [2] * 3
self.cropped_inc = [0, 96, 0, 0]
self.heatmap_size = self.image_size
self.sigma_regularization = 100
self.sigma_scale = 1000.0
self.cropped_training = True
self.output_folder = os.path.join('./output/vertebrae_localization/',
network.__name__, unet.__name__,
output_folder_name, str(cv),
self.output_folder_timestamp())
dataset_parameters = {
'base_folder': self.local_base_folder,
'image_size': self.image_size,
'image_spacing': self.image_spacing,
'cv': cv,
'input_gaussian_sigma': 0.75,
'generate_landmarks': True,
'generate_landmark_mask': True,
'translate_to_center_landmarks': True,
'translate_by_random_factor': True,
'save_debug_images': self.save_debug_images
}
dataset = Dataset(**dataset_parameters)
if self.use_pyro_dataset:
server_name = '@localhost:51232'
uri = 'PYRO:verse_dataset' + server_name
print('using pyro uri', uri)
self.dataset_train = PyroClientDataset(uri, **dataset_parameters)
else:
self.dataset_train = dataset.dataset_train()
self.dataset_val = dataset.dataset_val()
self.point_statistics_names = [
'pe_mean', 'pe_stdev', 'pe_median', 'num_correct'
]
self.additional_summaries_placeholders_val = dict([
(name, create_summary_placeholder(name))
for name in self.point_statistics_names
])
def loss_function(self, pred, target, mask=None):
"""
L2 loss function calculated with prediction and target.
:param pred: The predicted image.
:param target: The target image.
:param mask: If not none, calculate loss only pixels, where mask == 1
:return: L2 loss of (pred - target) / batch_size
"""
batch_size, _, _ = get_batch_channel_image_size(pred, self.data_format)
if mask is not None:
return tf.nn.l2_loss((pred - target) * mask) / batch_size
else:
return tf.nn.l2_loss(pred - target) / batch_size
def loss_function_sigmas(self, sigmas, valid_landmarks):
"""
L2 loss function for sigmas. Only calculated for values ver valid_landmarks == 1.
:param sigmas: Sigma variables.
:param valid_landmarks: Valid landmarks. Needs to have same shape as sigmas.
:return: L2 loss of sigmas * valid_landmarks.
"""
return self.sigma_regularization * tf.nn.l2_loss(
sigmas * valid_landmarks)
def init_networks(self):
"""
Initialize networks and placeholders.
"""
network_image_size = list(reversed(self.image_size))
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image', [1] + network_image_size),
('landmarks', [self.num_landmarks, 4]),
('landmark_mask', [1] + network_image_size)
])
else:
data_generator_entries = OrderedDict([
('image', network_image_size + [1]),
('landmarks', [self.num_landmarks, 4]),
('landmark_mask', network_image_size + [1])
])
data_generator_types = {'image': tf.float32}
# create sigmas variable
sigmas = tf.get_variable('sigmas', [self.num_landmarks],
initializer=tf.constant_initializer(
self.heatmap_sigma))
if not self.learnable_sigma:
sigmas = tf.stop_gradient(sigmas)
mean_sigmas = tf.reduce_mean(sigmas)
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(
coord=self.coord,
dataset=self.dataset_train,
data_names_and_shapes=data_generator_entries,
data_types=data_generator_types,
batch_size=self.batch_size)
data, target_landmarks, landmark_mask = self.train_queue.dequeue()
target_heatmaps = generate_heatmap_target(list(
reversed(self.heatmap_size)),
target_landmarks,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
prediction, local_prediction, spatial_prediction = training_net(
data,
num_labels=self.num_landmarks,
is_training=True,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
# losses
self.loss_net = self.loss_function(target=target_heatmaps,
pred=prediction,
mask=landmark_mask)
self.loss_sigmas = self.loss_function_sigmas(sigmas,
target_landmarks[0, :, 0])
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss = self.loss_net + self.loss_reg + self.loss_sigmas
# solver
global_step = tf.Variable(self.current_iter, trainable=False)
learning_rate = tf.train.piecewise_constant(
global_step, self.learning_rate_boundaries, self.learning_rates)
optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate,
momentum=0.99,
use_nesterov=True)
unclipped_gradients, variables = zip(
*optimizer.compute_gradients(self.loss))
norm = tf.global_norm(unclipped_gradients)
if self.clip_gradient_global_norm > 0:
gradients, _ = tf.clip_by_global_norm(
unclipped_gradients, self.clip_gradient_global_norm)
else:
gradients = unclipped_gradients
self.optimizer = optimizer.apply_gradients(zip(gradients, variables),
global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg),
('loss_sigmas', self.loss_sigmas),
('mean_sigmas', mean_sigmas),
('gradient_norm', norm)])
# build val graph
self.data_val, self.target_landmarks_val, self.landmark_mask_val = create_placeholders_tuple(
data_generator_entries,
data_types=data_generator_types,
shape_prefix=[1])
self.target_heatmaps_val = generate_heatmap_target(
list(reversed(self.heatmap_size)),
self.target_landmarks_val,
sigmas,
scale=self.sigma_scale,
normalize=True,
data_format=self.data_format)
self.prediction_val, self.local_prediction_val, self.spatial_prediction_val = training_net(
self.data_val,
num_labels=self.num_landmarks,
is_training=False,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
if self.has_validation_groundtruth:
self.loss_val = self.loss_function(target=self.target_heatmaps_val,
pred=self.prediction_val)
self.val_losses = OrderedDict([
('loss', self.loss_val), ('loss_reg', self.loss_reg),
('loss_sigmas', tf.constant(0, tf.float32)),
('mean_sigmas', tf.constant(0, tf.float32)),
('gradient_norm', tf.constant(0, tf.float32))
])
def test_cropped_image(self, dataset_entry):
"""
Perform inference on a dataset_entry with the validation network. Performs cropped prediction and merges outputs as maxima.
:param dataset_entry: A dataset entry from the dataset.
:return: input image (np.array), target heatmaps (np.array), predicted heatmaps, transformation (sitk.Transform)
"""
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
transformation = transformations['image']
image_size_np = [1] + list(reversed(self.image_size))
labels_size_np = [self.num_landmarks] + list(reversed(self.image_size))
full_image = generators['image']
landmarks = generators['landmarks']
image_tiler = ImageTiler(full_image.shape, image_size_np,
self.cropped_inc, True, -1)
landmark_tiler = LandmarkTiler(full_image.shape, image_size_np,
self.cropped_inc)
prediction_tiler = ImageTiler(
(self.num_landmarks, ) + full_image.shape[1:], labels_size_np,
self.cropped_inc, True, 0)
for image_tiler, landmark_tiler, prediction_tiler in zip(
image_tiler, landmark_tiler, prediction_tiler):
current_image = image_tiler.get_current_data(full_image)
current_landmarks = landmark_tiler.get_current_data(landmarks)
if self.has_validation_groundtruth:
feed_dict = {
self.data_val:
np.expand_dims(current_image, axis=0),
self.target_landmarks_val:
np.expand_dims(current_landmarks, axis=0)
}
run_tuple = self.sess.run(
(self.prediction_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(),
feed_dict=feed_dict)
else:
feed_dict = {
self.data_val: np.expand_dims(current_image, axis=0)
}
run_tuple = self.sess.run((self.prediction_val, ),
feed_dict=feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
image_tiler.set_current_data(current_image)
prediction_tiler.set_current_data(prediction)
return image_tiler.output_image, prediction_tiler.output_image, transformation
def test(self):
"""
The test function. Performs inference on the the validation images and calculates the loss.
"""
print('Testing...')
channel_axis = 0
if self.data_format == 'channels_last':
channel_axis = 3
if self.use_spine_postprocessing:
heatmap_maxima = HeatmapTest(channel_axis,
False,
return_multiple_maxima=True,
min_max_distance=7,
min_max_value=0.25,
multiple_min_max_value_factor=0.1)
spine_postprocessing = SpinePostprocessing(
num_landmarks=self.num_landmarks,
image_spacing=self.image_spacing)
else:
heatmap_maxima = HeatmapTest(channel_axis, False)
landmark_statistics = LandmarkStatistics()
landmarks = {}
num_entries = self.dataset_val.num_entries()
for i in range(num_entries):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
input_image = datasources['image']
target_landmarks = datasources['landmarks']
image, prediction, transformation = self.test_cropped_image(
dataset_entry)
if self.save_output_images:
if self.save_output_images_as_uint:
image_normalization = 'min_max'
heatmap_normalization = (0, 1)
output_image_type = np.uint8
else:
image_normalization = None
heatmap_normalization = None
output_image_type = np.float32
origin = transformation.TransformPoint(np.zeros(3, np.float64))
utils.io.image.write_multichannel_np(
image,
self.output_file_for_current_iteration(current_id +
'_input.mha'),
output_normalization_mode=image_normalization,
data_format=self.data_format,
image_type=output_image_type,
spacing=self.image_spacing,
origin=origin)
utils.io.image.write_multichannel_np(
prediction,
self.output_file_for_current_iteration(current_id +
'_prediction.mha'),
output_normalization_mode=heatmap_normalization,
data_format=self.data_format,
image_type=output_image_type,
spacing=self.image_spacing,
origin=origin)
if self.use_spine_postprocessing:
local_maxima_landmarks = heatmap_maxima.get_landmarks(
prediction, input_image, self.image_spacing,
transformation)
landmark_sequence = spine_postprocessing.postprocess_landmarks(
local_maxima_landmarks, prediction.shape)
landmarks[current_id] = landmark_sequence
else:
maxima_landmarks = heatmap_maxima.get_landmarks(
prediction, input_image, self.image_spacing,
transformation)
landmarks[current_id] = maxima_landmarks
if self.has_validation_groundtruth:
landmark_statistics.add_landmarks(current_id,
landmark_sequence,
target_landmarks)
print_progress_bar(i,
num_entries,
prefix='Testing ',
suffix=' complete')
utils.io.landmark.save_points_csv(
landmarks, self.output_file_for_current_iteration('points.csv'))
# finalize loss values
if self.has_validation_groundtruth:
print(landmark_statistics.get_pe_overview_string())
print(landmark_statistics.get_correct_id_string(20.0))
summary_values = OrderedDict(
zip(
self.point_statistics_names,
list(landmark_statistics.get_pe_statistics()) +
[landmark_statistics.get_correct_id(20)]))
# finalize loss values
self.val_loss_aggregator.finalize(self.current_iter,
summary_values)
overview_string = landmark_statistics.get_overview_string(
[2, 2.5, 3, 4, 10, 20], 10, 20.0)
utils.io.text.save_string_txt(
overview_string,
self.output_file_for_current_iteration('eval.txt'))
class MainLoop(MainLoopBase):
def __init__(self,
cv,
network,
unet,
network_parameters,
learning_rate,
output_folder_name=''):
"""
Initializer.
:param cv: The cv fold. 0, 1, 2 for CV; 'train_all' for training on whole dataset.
:param network: The used network. Usually network_u.
:param unet: The specific instance of the U-Net. Usually UnetClassicAvgLinear3d.
:param network_parameters: The network parameters passed to unet.
:param learning_rate: The initial learning rate.
:param output_folder_name: The output folder name that is used for distinguishing experiments.
"""
super().__init__()
self.batch_size = 1
self.learning_rates = [
learning_rate, learning_rate * 0.5, learning_rate * 0.1
]
self.learning_rate_boundaries = [20000, 30000]
self.max_iter = 50000
self.test_iter = 5000
self.disp_iter = 100
self.snapshot_iter = 5000
self.test_initialization = False
self.current_iter = 0
self.reg_constant = 0.000001
self.num_labels = 1
self.num_labels_all = 26
self.data_format = 'channels_first'
self.channel_axis = 1
self.network = network
self.unet = unet
self.network_parameters = network_parameters
self.padding = 'same'
self.clip_gradient_global_norm = 1.0
self.use_pyro_dataset = False
self.save_output_images = True
self.save_output_images_as_uint = True # set to False, if you want to see the direct network output
self.save_debug_images = False
self.has_validation_groundtruth = cv in [0, 1, 2]
self.local_base_folder = '../verse2019_dataset'
self.image_size = [128, 128, 96]
self.image_spacing = [1] * 3
self.output_folder = os.path.join('./output/vertebrae_segmentation/',
network.__name__, unet.__name__,
output_folder_name, str(cv),
self.output_folder_timestamp())
dataset_parameters = {
'base_folder': self.local_base_folder,
'image_size': self.image_size,
'image_spacing': self.image_spacing,
'cv': cv,
'input_gaussian_sigma': 0.75,
'label_gaussian_sigma': 1.0,
'heatmap_sigma': 3.0,
'generate_single_vertebrae_heatmap': True,
'generate_single_vertebrae': True,
'save_debug_images': self.save_debug_images
}
dataset = Dataset(**dataset_parameters)
if self.use_pyro_dataset:
server_name = '@localhost:51232'
uri = 'PYRO:verse_dataset' + server_name
print('using pyro uri', uri)
self.dataset_train = PyroClientDataset(uri, **dataset_parameters)
else:
self.dataset_train = dataset.dataset_train()
self.dataset_val = dataset.dataset_val()
self.dice_names = ['mean_dice'] + list(
map(lambda x: 'dice_{}'.format(x), range(self.num_labels_all)))
self.hausdorff_names = ['mean_h'] + list(
map(lambda x: 'h_{}'.format(x), range(self.num_labels)))
self.additional_summaries_placeholders_val = dict([
(name, create_summary_placeholder(name))
for name in (self.dice_names + self.hausdorff_names)
])
self.loss_function = sigmoid_cross_entropy_with_logits
self.setup_base_folder = os.path.join(self.local_base_folder, 'setup')
if cv in [0, 1, 2]:
self.cv_folder = os.path.join(self.setup_base_folder,
os.path.join('cv', str(cv)))
self.test_file = os.path.join(self.cv_folder, 'val.txt')
else:
self.test_file = os.path.join(self.setup_base_folder,
'train_all.txt')
self.valid_landmarks_file = os.path.join(self.setup_base_folder,
'valid_landmarks.csv')
self.test_id_list = utils.io.text.load_list(self.test_file)
self.valid_landmarks = utils.io.text.load_dict_csv(
self.valid_landmarks_file)
def init_networks(self):
"""
Initialize networks and placeholders.
"""
network_image_size = list(reversed(self.image_size))
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image', [1] + network_image_size),
('single_label', [self.num_labels] + network_image_size),
('single_heatmap', [1] + network_image_size)
])
else:
data_generator_entries = OrderedDict([
('image', network_image_size + [1]),
('single_label', network_image_size + [self.num_labels]),
('single_heatmap', network_image_size + [1])
])
data_generator_types = {'image': tf.float32, 'labels': tf.uint8}
# create model with shared weights between train and val
training_net = tf.make_template('net', self.network)
# build train graph
self.train_queue = DataGenerator(
coord=self.coord,
dataset=self.dataset_train,
data_names_and_shapes=data_generator_entries,
data_types=data_generator_types,
batch_size=self.batch_size)
data, mask, single_heatmap = self.train_queue.dequeue()
data_heatmap_concat = tf.concat([data, single_heatmap], axis=1)
prediction = training_net(data_heatmap_concat,
num_labels=self.num_labels,
is_training=True,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
# losses
self.loss_net = self.loss_function(labels=mask,
logits=prediction,
data_format=self.data_format)
self.loss_reg = get_reg_loss(self.reg_constant)
self.loss = self.loss_net + self.loss_reg
# solver
global_step = tf.Variable(self.current_iter, trainable=False)
learning_rate = tf.train.piecewise_constant(
global_step, self.learning_rate_boundaries, self.learning_rates)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
unclipped_gradients, variables = zip(
*optimizer.compute_gradients(self.loss))
norm = tf.global_norm(unclipped_gradients)
if self.clip_gradient_global_norm > 0:
gradients, _ = tf.clip_by_global_norm(
unclipped_gradients, self.clip_gradient_global_norm)
else:
gradients = unclipped_gradients
self.optimizer = optimizer.apply_gradients(zip(gradients, variables),
global_step=global_step)
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg),
('gradient_norm', norm)])
# build val graph
self.data_val, self.mask_val, self.single_heatmap_val = create_placeholders_tuple(
data_generator_entries,
data_types=data_generator_types,
shape_prefix=[1])
self.data_heatmap_concat_val = tf.concat(
[self.data_val, self.single_heatmap_val], axis=1)
self.prediction_val = training_net(self.data_heatmap_concat_val,
num_labels=self.num_labels,
is_training=False,
actual_network=self.unet,
padding=self.padding,
**self.network_parameters)
self.prediction_softmax_val = tf.nn.sigmoid(self.prediction_val)
if self.has_validation_groundtruth:
self.loss_val = self.loss_function(labels=self.mask_val,
logits=self.prediction_val,
data_format=self.data_format)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg),
('gradient_norm',
tf.constant(0, tf.float32))])
def test_full_image(self, dataset_entry):
"""
Perform inference on a dataset_entry with the validation network.
:param dataset_entry: A dataset entry from the dataset.
:return: input image (np.array), network prediction (np.array), transformation (sitk.Transform)
"""
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
if self.has_validation_groundtruth:
feed_dict = {
self.data_val:
np.expand_dims(generators['image'], axis=0),
self.mask_val:
np.expand_dims(generators['single_label'], axis=0),
self.single_heatmap_val:
np.expand_dims(generators['single_heatmap'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run(
(self.prediction_softmax_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(),
feed_dict=feed_dict)
else:
feed_dict = {
self.data_val:
np.expand_dims(generators['image'], axis=0),
self.single_heatmap_val:
np.expand_dims(generators['single_heatmap'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run((self.prediction_softmax_val, ),
feed_dict=feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
transformation = transformations['image']
image = generators['image']
return image, prediction, transformation
def test(self):
"""
The test function. Performs inference on the the validation images and calculates the loss.
"""
print('Testing...')
channel_axis = 0
if self.data_format == 'channels_last':
channel_axis = 3
labels = list(range(self.num_labels_all))
segmentation_test = SegmentationTest(labels,
channel_axis=channel_axis,
largest_connected_component=False,
all_labels_are_connected=False)
segmentation_statistics = SegmentationStatistics(
list(range(self.num_labels_all)),
self.output_folder_for_current_iteration(),
metrics=OrderedDict([('dice', DiceMetric()),
('h', HausdorffDistanceMetric())]))
filter_largest_cc = True
# iterate over all images
for i, image_id in enumerate(self.test_id_list):
first = True
prediction_resampled_np = None
input_image = None
groundtruth = None
# iterate over all valid landmarks
for landmark_id in self.valid_landmarks[image_id]:
dataset_entry = self.dataset_val.get({
'image_id': image_id,
'landmark_id': landmark_id
})
datasources = dataset_entry['datasources']
if first:
input_image = datasources['image']
if self.has_validation_groundtruth:
groundtruth = datasources['labels']
prediction_resampled_np = np.zeros(
[self.num_labels_all] +
list(reversed(input_image.GetSize())),
dtype=np.float16)
prediction_resampled_np[0, ...] = 0.5
first = False
image, prediction, transformation = self.test_full_image(
dataset_entry)
if filter_largest_cc:
prediction_thresh_np = (prediction > 0.5).astype(np.uint8)
largest_connected_component = utils.np_image.largest_connected_component(
prediction_thresh_np[0])
prediction_thresh_np[largest_connected_component[None, ...]
== 1] = 0
prediction[prediction_thresh_np == 1] = 0
if self.save_output_images:
if self.save_output_images_as_uint:
image_normalization = 'min_max'
label_normalization = (0, 1)
output_image_type = np.uint8
else:
image_normalization = None
label_normalization = None
output_image_type = np.float32
origin = transformation.TransformPoint(
np.zeros(3, np.float64))
utils.io.image.write_multichannel_np(
image,
self.output_file_for_current_iteration(image_id + '_' +
landmark_id +
'_input.mha'),
output_normalization_mode=image_normalization,
data_format=self.data_format,
image_type=output_image_type,
spacing=self.image_spacing,
origin=origin)
utils.io.image.write_multichannel_np(
prediction,
self.output_file_for_current_iteration(
image_id + '_' + landmark_id + '_prediction.mha'),
output_normalization_mode=label_normalization,
data_format=self.data_format,
image_type=output_image_type,
spacing=self.image_spacing,
origin=origin)
prediction_resampled_sitk = utils.sitk_image.transform_np_output_to_sitk_input(
output_image=prediction,
output_spacing=self.image_spacing,
channel_axis=channel_axis,
input_image_sitk=input_image,
transform=transformation,
interpolator='linear',
output_pixel_type=sitk.sitkFloat32)
#utils.io.image.write(prediction_resampled_sitk[0], self.output_file_for_current_iteration(image_id + '_' + landmark_id + '_resampled.mha'))
if self.data_format == 'channels_first':
prediction_resampled_np[int(landmark_id) + 1,
...] = utils.sitk_np.sitk_to_np(
prediction_resampled_sitk[0])
else:
prediction_resampled_np[..., int(landmark_id) +
1] = utils.sitk_np.sitk_to_np(
prediction_resampled_sitk[0])
prediction_labels = segmentation_test.get_label_image(
prediction_resampled_np, reference_sitk=input_image)
# delete to save memory
del prediction_resampled_np
utils.io.image.write(
prediction_labels,
self.output_file_for_current_iteration(image_id + '.mha'))
if self.has_validation_groundtruth:
segmentation_statistics.add_labels(image_id, prediction_labels,
groundtruth)
print_progress_bar(i,
len(self.test_id_list),
prefix='Testing ',
suffix=' complete')
# finalize loss values
if self.has_validation_groundtruth:
segmentation_statistics.finalize()
dice_list = segmentation_statistics.get_metric_mean_list('dice')
mean_dice = np.nanmean(dice_list)
dice_list = [mean_dice] + dice_list
hausdorff_list = segmentation_statistics.get_metric_mean_list('h')
mean_hausdorff = np.mean(hausdorff_list)
hausdorff_list = [mean_hausdorff] + hausdorff_list
summary_values = OrderedDict(
list(zip(self.dice_names, dice_list)) +
list(zip(self.hausdorff_names, hausdorff_list)))
self.val_loss_aggregator.finalize(self.current_iter,
summary_values=summary_values)
class MainLoop(MainLoopBase):
def __init__(self, cv, network_id):
super().__init__()
self.cv = cv
self.network_id = network_id
self.output_folder = '{}_cv{}'.format(
network_id, cv) + '/' + self.output_folder_timestamp()
self.batch_size = 8
self.learning_rate = 0.0000001
self.max_iter = 20000
self.test_iter = 5000
self.disp_iter = 100
self.snapshot_iter = self.test_iter
self.test_initialization = True
self.current_iter = 0
self.reg_constant = 0.00005
self.invert_transformation = False
image_sizes = {
'scn': [256, 256],
'unet': [256, 256],
'downsampling': [256, 256],
'conv': [128, 128],
'scn_mmwhs': [256, 256]
}
heatmap_sizes = {
'scn': [256, 256],
'unet': [256, 256],
'downsampling': [64, 64],
'conv': [128, 128],
'scn_mmwhs': [256, 256]
}
sigmas = {
'scn': 3.0,
'unet': 3.0,
'downsampling': 1.5,
'conv': 3.0,
'scn_mmwhs': 3.0
}
self.image_size = image_sizes[self.network_id]
self.heatmap_size = heatmap_sizes[self.network_id]
self.sigma = sigmas[self.network_id]
self.image_channels = 1
self.num_landmarks = 37
self.data_format = 'channels_first'
self.save_debug_images = False
self.base_folder = 'hand_xray_dataset/'
dataset = Dataset(self.image_size, self.heatmap_size,
self.num_landmarks, self.sigma, self.base_folder,
self.cv, self.data_format, self.save_debug_images)
self.dataset_train = dataset.dataset_train()
self.dataset_train.get_next()
self.dataset_val = dataset.dataset_val()
networks = {
'scn': network_scn,
'unet': network_unet,
'downsampling': network_downsampling,
'conv': network_conv,
'scn_mmwhs': network_scn_mmwhs
}
self.network = networks[self.network_id]
self.loss_function = lambda x, y: tf.nn.l2_loss(
x - y) / get_batch_channel_image_size(x, self.data_format)[0]
self.files_to_copy = ['main.py', 'network.py', 'dataset.py']
def initNetworks(self):
net = tf.make_template('net', self.network)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks',
[self.num_landmarks] + list(reversed(self.heatmap_size)))
])
data_generator_entries_val = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks',
[self.num_landmarks] + list(reversed(self.heatmap_size)))
])
else:
raise NotImplementedError
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size,
n_threads=8)
placeholders = self.train_queue.dequeue()
image = placeholders[0]
landmarks = placeholders[1]
prediction = net(image,
num_landmarks=self.num_landmarks,
is_training=True,
data_format=self.data_format)
self.loss_net = self.loss_function(landmarks, prediction)
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
if self.reg_constant > 0:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg
else:
self.loss_reg = 0
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg)])
#self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss)
self.optimizer = tf.train.MomentumOptimizer(
learning_rate=self.learning_rate, momentum=0.99,
use_nesterov=True).minimize(self.loss)
# build val graph
self.val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(
data_generator_entries_val, shape_prefix=[1])
self.image_val = self.val_placeholders['image']
self.landmarks_val = self.val_placeholders['landmarks']
self.prediction_val = net(self.image_val,
num_landmarks=self.num_landmarks,
is_training=False,
data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.landmarks_val,
self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg)])
def test_full_image(self, dataset_entry):
generators = dataset_entry['generators']
transformations = dataset_entry['transformations']
feed_dict = {
self.val_placeholders['image']:
np.expand_dims(generators['image'], axis=0),
self.val_placeholders['landmarks']:
np.expand_dims(generators['landmarks'], axis=0)
}
# run loss and update loss accumulators
run_tuple = self.sess.run((self.prediction_val, self.loss_val) +
self.val_loss_aggregator.get_update_ops(),
feed_dict)
prediction = np.squeeze(run_tuple[0], axis=0)
image = generators['image']
transformation = transformations['image']
return image, prediction, transformation
def test(self):
heatmap_test = HeatmapTest(channel_axis=0, invert_transformation=False)
landmark_statistics = LandmarkStatistics()
landmarks = {}
for i in range(self.dataset_val.num_entries()):
dataset_entry = self.dataset_val.get_next()
current_id = dataset_entry['id']['image_id']
datasources = dataset_entry['datasources']
reference_image = datasources['image_datasource']
groundtruth_landmarks = datasources['landmarks_datasource']
image, prediction, transform = self.test_full_image(dataset_entry)
utils.io.image.write_np(
(prediction * 128).astype(np.int8),
self.output_file_for_current_iteration(current_id +
'_heatmap.mha'))
predicted_landmarks = heatmap_test.get_landmarks(
prediction, reference_image, transformation=transform)
tensorflow_train.utils.tensorflow_util.print_progress_bar(
i, self.dataset_val.num_entries())
landmarks[current_id] = predicted_landmarks
landmark_statistics.add_landmarks(current_id,
predicted_landmarks,
groundtruth_landmarks,
normalization_factor=50,
normalization_indizes=[1, 5])
tensorflow_train.utils.tensorflow_util.print_progress_bar(
self.dataset_val.num_entries(), self.dataset_val.num_entries())
print('ipe', landmark_statistics.get_ipe_statistics())
print('pe', landmark_statistics.get_pe_statistics())
print('outliers',
landmark_statistics.get_num_outliers([2.0, 4.0, 10.0]))
# finalize loss values
self.val_loss_aggregator.finalize(self.current_iter)
utils.io.landmark.save_points_csv(
landmarks,
self.output_file_for_current_iteration('prediction.csv'))