Modification of the definition of networks

dh_segment/estimator_fn.py

@@ -2,7 +2,6 @@
 from .utils import PredictionType, ModelParams, TrainingParams, \
     class_to_label_image, multiclass_to_label_image
 import numpy as np
-from .network.model import inference_resnet_v1_50, inference_vgg16, inference_u_net
 
 
 def model_fn(mode, features, labels, params):
@@ -18,45 +17,22 @@ def model_fn(mode, features, labels, params):
         input_images = tf.pad(input_images, [[0, 0], [margin, margin], [margin, margin], [0, 0]],
                               mode='SYMMETRIC', name='mirror_padding')
 
-    if model_params.pretrained_model_name == 'vgg16':
-        network_output = inference_vgg16(input_images,
-                                         model_params,
-                                         model_params.n_classes,
-                                         use_batch_norm=model_params.batch_norm,
-                                         weight_decay=model_params.weight_decay,
-                                         is_training=(mode == tf.estimator.ModeKeys.TRAIN)
-                                         )
-        key_restore_model = 'vgg_16'
+    encoder_class = model_params.get_encoder()
+    encoder = encoder_class(**model_params.encoder_params)
+    decoder_class = model_params.get_decoder()
+    decoder = decoder_class(**model_params.decoder_params)
 
-    elif model_params.pretrained_model_name == 'resnet50':
-        network_output = inference_resnet_v1_50(input_images,
-                                                model_params,
-                                                model_params.n_classes,
-                                                use_batch_norm=model_params.batch_norm,
-                                                weight_decay=model_params.weight_decay,
-                                                is_training=(mode == tf.estimator.ModeKeys.TRAIN)
-                                                )
-        key_restore_model = 'resnet_v1_50'
-    elif model_params.pretrained_model_name == 'unet':
-        network_output = inference_u_net(input_images,
-                                         model_params,
-                                         model_params.n_classes,
-                                         use_batch_norm=model_params.batch_norm,
-                                         weight_decay=model_params.weight_decay,
-                                         is_training=(mode == tf.estimator.ModeKeys.TRAIN)
-                                         )
-        key_restore_model = None
-    else:
-        raise NotImplementedError
+    feature_maps = encoder(input_images)
+    network_output = decoder(feature_maps, num_classes=model_params.n_classes)
 
     if mode == tf.estimator.ModeKeys.TRAIN:
-        if key_restore_model is not None:
+        pretrained_file, pretrained_vars = encoder.pretrained_information()
+        if pretrained_file:
             # Pretrained weights as initialization
-            pretrained_restorer = tf.train.Saver(var_list=[v for v in tf.global_variables()
-                                                           if key_restore_model in v.name])
+            pretrained_restorer = tf.train.Saver(var_list=pretrained_vars)
 
             def init_fn(scaffold, session):
-                pretrained_restorer.restore(session, model_params.pretrained_model_file)
+                pretrained_restorer.restore(session, pretrained_file)
         else:
             init_fn = None
     else:
@@ -92,8 +68,10 @@ def init_fn(scaffold, session):
         if prediction_type == PredictionType.CLASSIFICATION:
             onehot_labels = tf.one_hot(indices=labels, depth=model_params.n_classes)
             with tf.name_scope("loss"):
-                per_pixel_loss = tf.nn.softmax_cross_entropy_with_logits(logits=network_output,
-                                                                         labels=onehot_labels, name='per_pixel_loss')
+                #per_pixel_loss = tf.nn.softmax_cross_entropy_with_logits(logits=network_output,
+                #                                                         labels=onehot_labels, name='per_pixel_loss')
+                per_pixel_loss = tf.nn.softmax_cross_entropy_with_logits_v2(logits=network_output,
+                                                                            labels=onehot_labels, name='per_pixel_loss')
                 if training_params.focal_loss_gamma > 0.0:
                     # Probability per pixel of getting the correct label
                     probs_correct_label = tf.reduce_max(tf.multiply(prediction_probs, onehot_labels))

dh_segment/io/input.py

@@ -17,9 +17,9 @@ class InputCase(Enum):
     INPUT_CSV = 'INPUT_CSV'
 
 
-def input_fn(input_data: Union[str, List[str]], params: dict, input_label_dir: str=None,
-             data_augmentation: bool=False, batch_size: int=5, make_patches: bool=False, num_epochs: int=1,
-             num_threads: int=4, image_summaries: bool=False):
+def input_fn(input_data: Union[str, List[str]], params: dict, input_label_dir: str = None,
+             data_augmentation: bool = False, batch_size: int = 5, make_patches: bool = False, num_epochs: int = 1,
+             num_threads: int = 4, image_summaries: bool = False, progressbar_description: str = 'Dataset'):
     """
     Input_fn for estimator
 
@@ -33,6 +33,7 @@ def input_fn(input_data: Union[str, List[str]], params: dict, input_label_dir: s
     :param num_epochs: number of epochs to cycle trough data (set it to None for infinite repeat)
     :param num_threads: number of thread to use in parallele when usin tf.data.Dataset.map
     :param image_summaries: boolean, whether to make tf.Summary to watch on tensorboard
+    :param progressbar_description: what will appear in the progressbar showing the number of files read
     :return: fn
     """
     training_params = utils.TrainingParams.from_dict(params['training_params'])
@@ -96,8 +97,9 @@ def _scaling_and_patch_fn(input_image, label_image):
 
     # Data augmentation
     def _augment_data_fn(input_image, label_image): \
-        return data_augmentation_fn(input_image, label_image, training_params.data_augmentation_flip_lr,
-                                    training_params.data_augmentation_flip_ud, training_params.data_augmentation_color)
+            return data_augmentation_fn(input_image, label_image, training_params.data_augmentation_flip_lr,
+                                        training_params.data_augmentation_flip_ud,
+                                        training_params.data_augmentation_color)
 
     # Assign color to class id
     def _assign_color_to_class_id(input_image, label_image):
@@ -112,27 +114,29 @@ def _assign_color_to_class_id(input_image, label_image):
             output['weight_maps'] = local_entropy(tf.equal(label_image, 1),
                                                   sigma=training_params.local_entropy_sigma)
         return output
+
     # ---
 
     # Finding the list of images to be used
     if isinstance(input_data, list):
         input_case = InputCase.INPUT_LIST
         input_image_filenames = input_data
-        print('Found {} images'.format(len(input_image_filenames)))
+        #print('Found {} images'.format(len(input_image_filenames)))
 
     elif os.path.isdir(input_data):
         input_case = InputCase.INPUT_DIR
         input_image_filenames = glob(os.path.join(input_data, '**', '*.jpg'),
                                      recursive=True) + \
                                 glob(os.path.join(input_data, '**', '*.png'),
                                      recursive=True)
-        print('Found {} images'.format(len(input_image_filenames)))
+        #print('Found {} images'.format(len(input_image_filenames)))
 
     elif os.path.isfile(input_data) and \
             input_data.endswith('.csv'):
         input_case = InputCase.INPUT_CSV
     else:
-        raise NotImplementedError('Input data should be a directory, a csv file or a list of filenames but got {}'.format(input_data))
+        raise NotImplementedError(
+            'Input data should be a directory, a csv file or a list of filenames but got {}'.format(input_data))
 
     # Finding the list of labelled images if available
     has_labelled_data = False
@@ -169,15 +173,16 @@ def _assign_color_to_class_id(input_image, label_image):
     def fn():
         if not has_labelled_data:
             encoded_filenames = [f.encode() for f in input_image_filenames]
-            dataset = tf.data.Dataset.from_generator(lambda: tqdm(encoded_filenames, desc='Dataset'),
+            dataset = tf.data.Dataset.from_generator(lambda: tqdm(encoded_filenames, desc=progressbar_description),
                                                      tf.string, tf.TensorShape([]))
             dataset = dataset.repeat(count=num_epochs)
             dataset = dataset.map(lambda filename: {'images': load_and_resize_image(filename, 3,
                                                                                     training_params.input_resized_size)})
         else:
             encoded_filenames = [(i.encode(), l.encode()) for i, l in zip(input_image_filenames, label_image_filenames)]
-            dataset = tf.data.Dataset.from_generator(lambda: tqdm(utils.shuffled(encoded_filenames), desc='Dataset'),
-                                                         (tf.string, tf.string), (tf.TensorShape([]), tf.TensorShape([])))
+            dataset = tf.data.Dataset.from_generator(lambda: tqdm(utils.shuffled(encoded_filenames),
+                                                                  desc=progressbar_description),
+                                                     (tf.string, tf.string), (tf.TensorShape([]), tf.TensorShape([])))
 
             dataset = dataset.repeat(count=num_epochs)
             dataset = dataset.map(_load_image_fn, num_threads).flat_map(_scaling_and_patch_fn)