Merge branch 'docstrings'

Author: qubvel

segmentation_models_pytorch/fpn/model.py

@@ -4,18 +4,35 @@
 
 
 class FPN(EncoderDecoder):
+    """FPN_ is a fully convolution neural network for image semantic segmentation
+    Args:
+        encoder_name: name of classification model (without last dense layers) used as feature
+                extractor to build segmentation model.
+        encoder_weights: one of ``None`` (random initialization), ``imagenet`` (pre-training on ImageNet).
+        decoder_pyramid_channels: a number of convolution filters in Feature Pyramid of FPN_.
+        decoder_segmentation_channels: a number of convolution filters in segmentation head of FPN_.
+        classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
+        dropout: spatial dropout rate in range (0, 1).
+        activation: one of [``sigmoid``, ``softmax``, None]
+
+    Returns:
+        ``keras.models.Model``: **FPN**
+
+    .. _FPN:
+        http://presentations.cocodataset.org/COCO17-Stuff-FAIR.pdf
+
+    """
 
     def __init__(
             self,
             encoder_name='resnet34',
             encoder_weights='imagenet',
             decoder_pyramid_channels=256,
-            decoder_segmenation_channels=128,
+            decoder_segmentation_channels=128,
             classes=1,
             dropout=0.2,
             activation='sigmoid',
     ):
-
         encoder = get_encoder(
             encoder_name,
             encoder_weights=encoder_weights
@@ -24,7 +41,7 @@ def __init__(
         decoder = FPNDecoder(
             encoder_channels=encoder.out_shapes,
             pyramid_channels=decoder_pyramid_channels,
-            segmentation_channels=decoder_segmenation_channels,
+            segmentation_channels=decoder_segmentation_channels,
             final_channels=classes,
             dropout=dropout,
         )

segmentation_models_pytorch/linknet/model.py

@@ -4,6 +4,26 @@
 
 
 class Linknet(EncoderDecoder):
+    """Linknet_ is a fully convolution neural network for fast image semantic segmentation
+
+    Note:
+        This implementation by default has 4 skip connections (original - 3).
+
+    Args:
+        encoder_name: name of classification model (without last dense layers) used as feature
+            extractor to build segmentation model.
+        encoder_weights: one of ``None`` (random initialization), ``imagenet`` (pre-training on ImageNet).
+        decoder_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
+            is used.
+        classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
+        activation: one of [``sigmoid``, ``softmax``, None]
+
+    Returns:
+        ``torch.nn.Module``: **Linknet**
+
+    .. _Linknet:
+        https://arxiv.org/pdf/1707.03718.pdf
+    """
 
     def __init__(
             self,
@@ -13,7 +33,6 @@ def __init__(
             classes=1,
             activation='sigmoid',
     ):
-
         encoder = get_encoder(
             encoder_name,
             encoder_weights=encoder_weights

segmentation_models_pytorch/pspnet/model.py

@@ -4,6 +4,28 @@
 
 
 class PSPNet(EncoderDecoder):
+    """PSPNet_ is a fully convolution neural network for image semantic segmentation
+
+    Args:
+        encoder_name: name of classification model used as feature
+                extractor to build segmentation model.
+        encoder_weights: one of ``None`` (random initialization), ``imagenet`` (pre-training on ImageNet).
+        psp_in_factor: one of 4, 8 and 16. Downsampling rate or in other words backbone depth
+            to construct PSP module on it.
+        psp_out_channels: number of filters in PSP block.
+        psp_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
+                is used.
+        psp_aux_output: if ``True`` add auxiliary classification output for encoder training
+        psp_dropout: spatial dropout rate between 0 and 1.
+        classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
+        activation: one of [``sigmoid``, ``softmax``, None]
+
+    Returns:
+        ``torch.nn.Module``: **PSPNet**
+
+    .. _PSPNet:
+        https://arxiv.org/pdf/1612.01105.pdf
+    """
 
     def __init__(
             self,
@@ -17,7 +39,6 @@ def __init__(
             dropout=0.2,
             activation='softmax',
     ):
-
         encoder = get_encoder(
             encoder_name,
             encoder_weights=encoder_weights

segmentation_models_pytorch/unet/model.py

@@ -4,6 +4,26 @@
 
 
 class Unet(EncoderDecoder):
+    """Unet_ is a fully convolution neural network for image semantic segmentation
+
+    Args:
+        encoder_name: name of classification model (without last dense layers) used as feature
+            extractor to build segmentation model.
+        encoder_weights: one of ``None`` (random initialization), ``imagenet`` (pre-training on ImageNet).
+        decoder_channels: list of numbers of ``Conv2D`` layer filters in decoder blocks
+        decoder_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
+            is used.
+        classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
+        activation: one of [``sigmoid``, ``softmax``, None]
+        center: if ``True`` add ``Conv2dReLU`` block on encoder head (useful for VGG models)
+
+    Returns:
+        ``torch.nn.Module``: **Unet**
+
+    .. _Unet:
+        https://arxiv.org/pdf/1505.04597
+
+    """
 
     def __init__(
             self,