Unet, PSPNet, Linknet docstrings

Author: qubvel

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 - ``(h, w, classes)``).
+        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 - ``(h, w, classes)``).
+        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,