upgrade contrib/segmentation_bf_aug

Author: wyli

niftynet/contrib/segmentation_bf_aug/segmentation_application_bfaug.py

@@ -1,86 +1,58 @@
 import tensorflow as tf
 
-from niftynet.application.base_application import BaseApplication
-from niftynet.engine.application_factory import \
-    ApplicationNetFactory, InitializerFactory, OptimiserFactory
-from niftynet.engine.application_variables import \
-    CONSOLE, NETWORK_OUTPUT, TF_SUMMARIES
-from niftynet.engine.sampler_grid import GridSampler
-from niftynet.engine.sampler_resize import ResizeSampler
-from niftynet.engine.sampler_uniform import UniformSampler
-from niftynet.engine.sampler_weighted import WeightedSampler
-from niftynet.engine.windows_aggregator_grid import GridSamplesAggregator
-from niftynet.engine.windows_aggregator_resize import ResizeSamplesAggregator
+from niftynet.application.segmentation_application import \
+    SegmentationApplication, SUPPORTED_INPUT
 from niftynet.io.image_reader import ImageReader
 from niftynet.layer.binary_masking import BinaryMaskingLayer
 from niftynet.layer.discrete_label_normalisation import \
     DiscreteLabelNormalisationLayer
 from niftynet.layer.histogram_normalisation import \
     HistogramNormalisationLayer
-from niftynet.layer.loss_segmentation import LossFunction
 from niftynet.layer.mean_variance_normalisation import \
     MeanVarNormalisationLayer
 from niftynet.layer.pad import PadLayer
-from niftynet.layer.post_processing import PostProcessingLayer
+from niftynet.layer.rand_bias_field import RandomBiasFieldLayer
 from niftynet.layer.rand_flip import RandomFlipLayer
 from niftynet.layer.rand_rotation import RandomRotationLayer
 from niftynet.layer.rand_spatial_scaling import RandomSpatialScalingLayer
-from niftynet.layer.rand_bias_field import RandomBiasFieldLayer
-
-SUPPORTED_INPUT = set(['image', 'label', 'weight', 'sampler'])
 
 
-class SegmentationApplicationBFAug(BaseApplication):
+class SegmentationApplicationBFAug(SegmentationApplication):
     REQUIRED_CONFIG_SECTION = "SEGMENTATION"
 
     def __init__(self, net_param, action_param, is_training):
-        super(SegmentationApplicationBFAug, self).__init__()
+        SegmentationApplication.__init__(
+            self, net_param, action_param, is_training)
         tf.logging.info('starting segmentation application')
-        self.is_training = is_training
-
-        self.net_param = net_param
-        self.action_param = action_param
-
-        self.data_param = None
-        self.segmentation_param = None
-        self.SUPPORTED_SAMPLING = {
-            'uniform': (self.initialise_uniform_sampler,
-                        self.initialise_grid_sampler,
-                        self.initialise_grid_aggregator),
-            'weighted': (self.initialise_weighted_sampler,
-                         self.initialise_grid_sampler,
-                         self.initialise_grid_aggregator),
-            'resize': (self.initialise_resize_sampler,
-                       self.initialise_resize_sampler,
-                       self.initialise_resize_aggregator),
-        }
 
     def initialise_dataset_loader(
             self, data_param=None, task_param=None, data_partitioner=None):
 
         self.data_param = data_param
         self.segmentation_param = task_param
 
-        # read each line of csv files into an instance of Subject
+        # initialise input image readers
         if self.is_training:
-            file_lists = []
-            if self.action_param.validation_every_n > 0:
-                file_lists.append(data_partitioner.train_files)
-                file_lists.append(data_partitioner.validation_files)
-            else:
-                file_lists.append(data_partitioner.train_files)
-
-            self.readers = []
-            for file_list in file_lists:
-                reader = ImageReader(SUPPORTED_INPUT)
-                reader.initialise(data_param, task_param, file_list)
-                self.readers.append(reader)
-
-        else:  # in the inference process use image input only
-            inference_reader = ImageReader(['image'])
-            file_list = data_partitioner.inference_files
-            inference_reader.initialise(data_param, task_param, file_list)
-            self.readers = [inference_reader]
+            reader_names = ('image', 'label', 'weight', 'sampler')
+        elif self.is_inference:
+            # in the inference process use `image` input only
+            reader_names = ('image',)
+        elif self.is_evaluation:
+            reader_names = ('image', 'label', 'inferred')
+        else:
+            tf.logging.fatal(
+                'Action `%s` not supported. Expected one of %s',
+                self.action, self.SUPPORTED_PHASES)
+            raise ValueError
+        try:
+            reader_phase = self.action_param.dataset_to_infer
+        except AttributeError:
+            reader_phase = None
+        file_lists = data_partitioner.get_file_lists_by(
+            phase=reader_phase, action=self.action)
+        self.readers = [
+            ImageReader(reader_names).initialise(
+                data_param, task_param, file_list) for file_list in file_lists]
 
         foreground_masking_layer = None
         if self.net_param.normalise_foreground_only:
@@ -148,209 +120,15 @@ def initialise_dataset_loader(
                     self.action_param.bias_field_range)
                 augmentation_layers.append(bias_field_layer)
 
-
         volume_padding_layer = []
         if self.net_param.volume_padding_size:
             volume_padding_layer.append(PadLayer(
                 image_name=SUPPORTED_INPUT,
                 border=self.net_param.volume_padding_size))
 
-        for reader in self.readers:
-            reader.add_preprocessing_layers(
-                volume_padding_layer +
-                normalisation_layers +
-                augmentation_layers)
-
-    def initialise_uniform_sampler(self):
-        self.sampler = [[UniformSampler(
-            reader=reader,
-            data_param=self.data_param,
-            batch_size=self.net_param.batch_size,
-            windows_per_image=self.action_param.sample_per_volume,
-            queue_length=self.net_param.queue_length) for reader in
-            self.readers]]
-
-    def initialise_weighted_sampler(self):
-        self.sampler = [[WeightedSampler(
-            reader=reader,
-            data_param=self.data_param,
-            batch_size=self.net_param.batch_size,
-            windows_per_image=self.action_param.sample_per_volume,
-            queue_length=self.net_param.queue_length) for reader in
-            self.readers]]
-
-    def initialise_resize_sampler(self):
-        self.sampler = [[ResizeSampler(
-            reader=reader,
-            data_param=self.data_param,
-            batch_size=self.net_param.batch_size,
-            shuffle_buffer=self.is_training,
-            queue_length=self.net_param.queue_length) for reader in
-            self.readers]]
+        self.readers[0].add_preprocessing_layers(
+            volume_padding_layer + normalisation_layers + augmentation_layers)
 
-    def initialise_grid_sampler(self):
-        self.sampler = [[GridSampler(
-            reader=reader,
-            data_param=self.data_param,
-            batch_size=self.net_param.batch_size,
-            spatial_window_size=self.action_param.spatial_window_size,
-            window_border=self.action_param.border,
-            queue_length=self.net_param.queue_length) for reader in
-            self.readers]]
-
-    def initialise_grid_aggregator(self):
-        self.output_decoder = GridSamplesAggregator(
-            image_reader=self.readers[0],
-            output_path=self.action_param.save_seg_dir,
-            window_border=self.action_param.border,
-            interp_order=self.action_param.output_interp_order)
-
-    def initialise_resize_aggregator(self):
-        self.output_decoder = ResizeSamplesAggregator(
-            image_reader=self.readers[0],
-            output_path=self.action_param.save_seg_dir,
-            window_border=self.action_param.border,
-            interp_order=self.action_param.output_interp_order)
-
-    def initialise_sampler(self):
-        if self.is_training:
-            self.SUPPORTED_SAMPLING[self.net_param.window_sampling][0]()
-        else:
-            self.SUPPORTED_SAMPLING[self.net_param.window_sampling][1]()
-
-    def initialise_network(self):
-        w_regularizer = None
-        b_regularizer = None
-        reg_type = self.net_param.reg_type.lower()
-        decay = self.net_param.decay
-        if reg_type == 'l2' and decay > 0:
-            from tensorflow.contrib.layers.python.layers import regularizers
-            w_regularizer = regularizers.l2_regularizer(decay)
-            b_regularizer = regularizers.l2_regularizer(decay)
-        elif reg_type == 'l1' and decay > 0:
-            from tensorflow.contrib.layers.python.layers import regularizers
-            w_regularizer = regularizers.l1_regularizer(decay)
-            b_regularizer = regularizers.l1_regularizer(decay)
-
-        self.net = ApplicationNetFactory.create(self.net_param.name)(
-            num_classes=self.segmentation_param.num_classes,
-            w_initializer=InitializerFactory.get_initializer(
-                name=self.net_param.weight_initializer),
-            b_initializer=InitializerFactory.get_initializer(
-                name=self.net_param.bias_initializer),
-            w_regularizer=w_regularizer,
-            b_regularizer=b_regularizer,
-            acti_func=self.net_param.activation_function)
-
-    def connect_data_and_network(self,
-                                 outputs_collector=None,
-                                 gradients_collector=None):
-        # def data_net(for_training):
-        #    with tf.name_scope('train' if for_training else 'validation'):
-        #        sampler = self.get_sampler()[0][0 if for_training else -1]
-        #        data_dict = sampler.pop_batch_op()
-        #        image = tf.cast(data_dict['image'], tf.float32)
-        #        return data_dict, self.net(image, is_training=for_training)
-
-        def switch_sampler(for_training):
-            with tf.name_scope('train' if for_training else 'validation'):
-                sampler = self.get_sampler()[0][0 if for_training else -1]
-                return sampler.pop_batch_op()
-
-        if self.is_training:
-            # if self.action_param.validation_every_n > 0:
-            #    data_dict, net_out = tf.cond(tf.logical_not(self.is_validation),
-            #                                 lambda: data_net(True),
-            #                                 lambda: data_net(False))
-            # else:
-            #    data_dict, net_out = data_net(True)
-            if self.action_param.validation_every_n > 0:
-                data_dict = tf.cond(tf.logical_not(self.is_validation),
-                                    lambda: switch_sampler(for_training=True),
-                                    lambda: switch_sampler(for_training=False))
-            else:
-                data_dict = switch_sampler(for_training=True)
-            image = tf.cast(data_dict['image'], tf.float32)
-            net_out = self.net(image, is_training=self.is_training)
-
-            with tf.name_scope('Optimiser'):
-                optimiser_class = OptimiserFactory.create(
-                    name=self.action_param.optimiser)
-                self.optimiser = optimiser_class.get_instance(
-                    learning_rate=self.action_param.lr)
-            loss_func = LossFunction(
-                n_class=self.segmentation_param.num_classes,
-                loss_type=self.action_param.loss_type)
-            data_loss = loss_func(
-                prediction=net_out,
-                ground_truth=data_dict.get('label', None),
-                weight_map=data_dict.get('weight', None))
-            reg_losses = tf.get_collection(
-                tf.GraphKeys.REGULARIZATION_LOSSES)
-            if self.net_param.decay > 0.0 and reg_losses:
-                reg_loss = tf.reduce_mean(
-                    [tf.reduce_mean(reg_loss) for reg_loss in reg_losses])
-                loss = data_loss + reg_loss
-            else:
-                loss = data_loss
-            grads = self.optimiser.compute_gradients(loss)
-            # collecting gradients variables
-            gradients_collector.add_to_collection([grads])
-            # collecting output variables
-            outputs_collector.add_to_collection(
-                var=data_loss, name='loss',
-                average_over_devices=False, collection=CONSOLE)
-            outputs_collector.add_to_collection(
-                var=data_loss, name='loss',
-                average_over_devices=True, summary_type='scalar',
-                collection=TF_SUMMARIES)
-
-            # outputs_collector.add_to_collection(
-            #    var=image*180.0, name='image',
-            #    average_over_devices=False, summary_type='image3_sagittal',
-            #    collection=TF_SUMMARIES)
-
-            # outputs_collector.add_to_collection(
-            #    var=image, name='image',
-            #    average_over_devices=False,
-            #    collection=NETWORK_OUTPUT)
-
-            # outputs_collector.add_to_collection(
-            #    var=tf.reduce_mean(image), name='mean_image',
-            #    average_over_devices=False, summary_type='scalar',
-            #    collection=CONSOLE)
-        else:
-            # converting logits into final output for
-            # classification probabilities or argmax classification labels
-            data_dict = switch_sampler(for_training=False)
-            image = tf.cast(data_dict['image'], tf.float32)
-            net_out = self.net(image, is_training=self.is_training)
-
-            output_prob = self.segmentation_param.output_prob
-            num_classes = self.segmentation_param.num_classes
-            if output_prob and num_classes > 1:
-                post_process_layer = PostProcessingLayer(
-                    'SOFTMAX', num_classes=num_classes)
-            elif not output_prob and num_classes > 1:
-                post_process_layer = PostProcessingLayer(
-                    'ARGMAX', num_classes=num_classes)
-            else:
-                post_process_layer = PostProcessingLayer(
-                    'IDENTITY', num_classes=num_classes)
-            net_out = post_process_layer(net_out)
-
-            outputs_collector.add_to_collection(
-                var=net_out, name='window',
-                average_over_devices=False, collection=NETWORK_OUTPUT)
-            outputs_collector.add_to_collection(
-                var=data_dict['image_location'], name='location',
-                average_over_devices=False, collection=NETWORK_OUTPUT)
-            init_aggregator = \
-                self.SUPPORTED_SAMPLING[self.net_param.window_sampling][2]
-            init_aggregator()
-
-    def interpret_output(self, batch_output):
-        if not self.is_training:
-            return self.output_decoder.decode_batch(
-                batch_output['window'], batch_output['location'])
-        return True
+        for reader in self.readers[1:]:
+            reader.add_preprocessing_layers(
+                volume_padding_layer + normalisation_layers)