Modification of the definition of networks

dh_segment/network/__init__.py

@@ -1,12 +1,14 @@
 _MODEL = [
-    'inference_vgg16',
-    'inference_resnet_v1_50',
-    'inference_u_net',
-    'vgg_16_fn',
-    'resnet_v1_50_fn'
+    'Encoder',
+    'Decoder',
+    'SimpleDecoder',
 ]
 
-__all__ = _MODEL
+_PRETRAINED = [
+    'ResnetV1_50',
+    'VGG16'
+]
+__all__ = _MODEL + _PRETRAINED
 
 from .model import *
 from .pretrained_models import *

dh_segment/network/model.py

@@ -13,15 +13,15 @@ def __call__(self, images: tf.Tensor) -> List[tf.Tensor]:
         """
 
         :param images: [NxHxWx3] float32 [0..255] input images
-        :return: a list of the feature maps in decreasing spatial resolution (first element is most likely the input
+        :return: a list of the feature maps in decreasing spatial resolution (first element is most likely the input \
         image itself, then the output of the first pooling op, etc...)
         """
         pass
 
     def pretrained_information(self) -> Tuple[Optional[str], Union[None, List, Dict]]:
         """
 
-        :return: The filename of the pretrained checkpoint and the corresponding variables (List of Dict mapping)
+        :return: The filename of the pretrained checkpoint and the corresponding variables (List of Dict mapping) \
         or `None` if no-pretraining is done
         """
         return None, None
@@ -32,14 +32,23 @@ class Decoder(ABC):
     def __call__(self, feature_maps: List[tf.Tensor], num_classes: int) -> tf.Tensor:
         """
 
-        :param feature_maps: list of feature maps, in decreasing spatial resolution, first one being at the original resolution
+        :param feature_maps: list of feature maps, in decreasing spatial resolution, first one being at the original \
+        resolution
         :return: [N,H,W,num_classes] float32 tensor of logit scores
         """
         pass
 
 
 class SimpleDecoder(Decoder):
-    def __init__(self, upsampling_dims: List[int], max_depth: int = None, train_batchnorm=False, weight_decay=0.):
+    """
+
+    :param upsampling_dims:
+    :param max_depth:
+    :param weight_decay:
+    :param self.batch_norm_fn:
+    """
+    def __init__(self, upsampling_dims: List[int], max_depth: int = None, train_batchnorm: bool=False,
+                 weight_decay: float=0.):
         self.upsampling_dims = upsampling_dims
         self.max_depth = max_depth
         self.weight_decay = weight_decay
@@ -105,6 +114,7 @@ def __call__(self, feature_maps: List[tf.Tensor], num_classes: int):
 def _get_image_shape_tensor(tensor: tf.Tensor) -> Union[Tuple[int, int], tf.Tensor]:
     """
     Get the image shape of the tensor
+
     :param tensor: Input image tensor [N,H,W,...]
     :return: a (int, int) tuple if shape is defined, otherwise the corresponding tf.Tensor value
     """
@@ -116,7 +126,7 @@ def _get_image_shape_tensor(tensor: tf.Tensor) -> Union[Tuple[int, int], tf.Tens
     return target_shape
 
 
-def _upsample_concat(pooled_layer: tf.Tensor, previous_layer: tf.Tensor, scope_name='UpsampleConcat'):
+def _upsample_concat(pooled_layer: tf.Tensor, previous_layer: tf.Tensor, scope_name: str='UpsampleConcat'):
     """
 
     :param pooled_layer: [N,H,W,C] coarse layer

dh_segment/network/pretrained_models.py

@@ -15,12 +15,22 @@ def mean_substraction(input_tensor, means=_VGG_MEANS):
 
 
 class ResnetV1_50(Encoder):
-    def __init__(self, train_batchnorm=False, blocks=4, weight_decay=0.0001,
-                    renorm=True, corrected_version=False):
+    """ResNet-50 implementation
+
+    :param train_batchnorm: Option to use batch norm
+    :param blocks: number of blocks (resnet blocks)
+    :param weight_decay: value of weight decay
+    :param batch_renorm: Option to use batch renorm
+    :param corrected_version: option to use the original resnet implementation (True) but less efficient than
+                              `slim`'s implementation
+    :param pretrained_file: path to the file (.ckpt) containing the pretrained weights
+    """
+    def __init__(self, train_batchnorm: bool=False, blocks: int=4, weight_decay: float=0.0001,
+                 batch_renorm: bool=True, corrected_version: bool=False):
         self.train_batchnorm = train_batchnorm
         self.blocks = blocks
         self.weight_decay = weight_decay
-        self.renorm = renorm
+        self.batch_renorm = batch_renorm
         self.corrected_version = corrected_version
         self.pretrained_file = os.path.join(get_data_folder(), 'resnet_v1_50.ckpt')
         if not os.path.exists(self.pretrained_file):
@@ -43,23 +53,21 @@ def __call__(self, images: tf.Tensor):
         outputs = []
 
         with slim.arg_scope(nets.resnet_v1.resnet_arg_scope(weight_decay=self.weight_decay, batch_norm_decay=0.999)), \
-             slim.arg_scope([layers.batch_norm], renorm_decay=0.95, renorm=self.renorm):
+             slim.arg_scope([layers.batch_norm], renorm_decay=0.95, renorm=self.batch_renorm):
             mean_substracted_tensor = mean_substraction(images)
             assert 0 < self.blocks <= 4
 
             if self.corrected_version:
-                def corrected_resnet_v1_block(scope, base_depth, num_units, stride):
-                    """Helper function for creating a resnet_v1 bottleneck block.
-
-                    Args:
-                      scope: The scope of the block.
-                      base_depth: The depth of the bottleneck layer for each unit.
-                      num_units: The number of units in the block.
-                      stride: The stride of the block, implemented as a stride in the last unit.
-                        All other units have stride=1.
-
-                    Returns:
-                      A resnet_v1 bottleneck block.
+                def corrected_resnet_v1_block(scope: str, base_depth: int, num_units: int, stride: int) -> tf.Tensor:
+                    """
+                    Helper function for creating a resnet_v1 bottleneck block.
+
+                    :param scope: The scope of the block.
+                    :param base_depth: The depth of the bottleneck layer for each unit.
+                    :param num_units: The number of units in the block.
+                    :param stride: The stride of the block, implemented as a stride in the last unit.
+                                   All other units have stride=1.
+                    :return: A resnet_v1 bottleneck block.
                     """
                     return nets.resnet_utils.Block(scope, nets.resnet_v1.bottleneck, [{
                         'depth': base_depth * 4,
@@ -119,7 +127,13 @@ def corrected_resnet_v1_block(scope, base_depth, num_units, stride):
 
 
 class VGG16(Encoder):
-    def __init__(self, blocks=5, weight_decay=0.0005):
+    """VGG-16 implementation
+
+    :param blocks: number of blocks (vgg blocks)
+    :param weight_decay: weight decay value
+    :param pretrained_file: path to the file (.ckpt) containing the pretrained weights
+    """
+    def __init__(self, blocks: int=5, weight_decay: float=0.0005):
         self.blocks = blocks
         self.weight_decay = weight_decay
         self.pretrained_file = os.path.join(get_data_folder(), 'vgg_16.ckpt')
@@ -140,36 +154,37 @@ def pretrained_information(self):
                                       and 'renorm' not in v.name]
 
     def __call__(self, images: tf.Tensor):
-        intermediate_levels = []
+        outputs = []
 
         with slim.arg_scope(nets.vgg.vgg_arg_scope(weight_decay=self.weight_decay)):
             with tf.variable_scope(None, 'vgg_16', [images]) as sc:
                 input_tensor = mean_substraction(images)
-                intermediate_levels.append(input_tensor)
+                outputs.append(input_tensor)
                 end_points_collection = sc.original_name_scope + '_end_points'
                 # Collect outputs for conv2d, fully_connected and max_pool2d.
                 with slim.arg_scope(
                         [layers.conv2d, layers.fully_connected, layers.max_pool2d],
                         outputs_collections=end_points_collection):
                     net = layers.repeat(
                         input_tensor, 2, layers.conv2d, 64, [3, 3], scope='conv1')
-                    intermediate_levels.append(net)
+                    outputs.append(net)
                     net = layers.max_pool2d(net, [2, 2], scope='pool1')
                     if self.blocks >= 2:
                         net = layers.repeat(net, 2, layers.conv2d, 128, [3, 3], scope='conv2')
-                        intermediate_levels.append(net)
+                        outputs.append(net)
                         net = layers.max_pool2d(net, [2, 2], scope='pool2')
                     if self.blocks >= 3:
                         net = layers.repeat(net, 3, layers.conv2d, 256, [3, 3], scope='conv3')
-                        intermediate_levels.append(net)
+                        outputs.append(net)
                         net = layers.max_pool2d(net, [2, 2], scope='pool3')
                     if self.blocks >= 4:
                         net = layers.repeat(net, 3, layers.conv2d, 512, [3, 3], scope='conv4')
-                        intermediate_levels.append(net)
+                        outputs.append(net)
                         net = layers.max_pool2d(net, [2, 2], scope='pool4')
                     if self.blocks >= 5:
                         net = layers.repeat(net, 3, layers.conv2d, 512, [3, 3], scope='conv5')
-                        intermediate_levels.append(net)
+                        outputs.append(net)
                         net = layers.max_pool2d(net, [2, 2], scope='pool5')
 
-                    return intermediate_levels
+                    # TODO : the output of the last max pool is not returned, shouldn't it be ?
+                    return outputs

dh_segment/utils/__init__.py

@@ -1,8 +1,5 @@
 _PARAMSCONFIG = [
     'PredictionType',
-    'VGG16ModelParams',
-    'ResNetModelParams',
-    'UNetModelParams',
     'ModelParams',
     'TrainingParams'
 ]

dh_segment/utils/misc.py

@@ -49,6 +49,7 @@ def get_class_from_name(full_class_name: str) -> Any:
     """
     Tries to load the class from its naming, will import the corresponding module.
     Raises an Error if it does not work.
+
     :param full_class_name: full name of the class, for instance `foo.bar.Baz`
     :return: the loaded class
     """
@@ -67,6 +68,12 @@ def get_data_folder() -> str:
 
 
 def download_file(url: str, output_file: str):
+    """
+
+    :param url:
+    :param output_file:
+    :return:
+    """
     def progress_hook(t):
         last_b = [0]
 

dh_segment/utils/params_config.py

@@ -13,7 +13,7 @@ class PredictionType:
     MULTILABEL = 'MULTILABEL'
 
     @classmethod
-    def parse(cls, prediction_type):
+    def parse(cls, prediction_type) -> 'PredictionType':
         if prediction_type == 'CLASSIFICATION':
             return PredictionType.CLASSIFICATION
         elif prediction_type == 'REGRESSION':
@@ -43,6 +43,14 @@ def check_params(self):
 
 
 class ModelParams(BaseParams):
+    """
+
+    :param encoder_name:
+    :param encoder_params:
+    :param decoder_name:
+    :param decoder_params:
+    :param n_classes:
+    """
     def __init__(self, **kwargs):
         self.encoder_name = kwargs.get('encoder_name', 'dh_segment.network.pretrained_models.ResnetV1_50')  # type: str
         self.encoder_params = kwargs.get('encoder_params', dict())  # type: dict