segmentation_models.pytorch

The segmentation_models_pytorch library provides several pre-defined segmentation architectures. You can instantiate these models using their respective classes. The supported architectures include:

• Unet
• UnetPlusPlus
• FPN
• PSPNet
• DeepLabV3
• DeepLabV3Plus
• Linknet
• MAnet
• PAN
• UPerNet
• Segformer
• DPT

SMP supports 12 different encoder-decoder model architectures. Key architectures include:

• Unet and Unet++
• MAnet
• Linknet
• FPN
• PSPNet
• PAN
• DeepLabV3 and DeepLabV3+
• UPerNet
• Segformer
• DPT

Encoders (also known as backbones) are used to extract intermediate features that are fed into the decoder for segmentation. SMP provides over 800+ pretrained convolution- and transform-based encoders, including support for timm encoders.

Choosing an encoder allows you to balance efficiency and accuracy:

• Lightweight encoders (e.g., mobilenet, mobileone) are suitable for low-latency or real-time inference on edge devices.
• High-capacity architectures (e.g., convnext, swin, mit) are better for complex tasks with many classes where superior accuracy is required.

When using 'Traditional-Style' Timm encoders, the models typically produce feature maps at specific downsampling scales relative to the input resolution. This is important for designing segmentation heads that expect multi-scale feature maps.

The standard scales are:

• 1/2
• 1/4
• 1/8
• 1/16
• 1/32

All segmentation models in segmentation_models_pytorch (SMP) are composed of four main components:

1. Encoder (backbone/feature extractor): Extracts features from the input image.
2. Decoder: A feature fusion block that uses encoder features to create the segmentation mask.
3. Segmentation head: The final head that reduces the number of channels from the decoder and upsamples the mask to match the input-output spatial resolution.
4. Classification head (optional): A head built on top of the deepest encoder features for auxiliary classification tasks.

After defining your MyEncoder class, you must register it in the smp.encoders.encoders dictionary so it can be used via the encoder_name parameter in model constructors.

Register the encoder with its pretrained_settings (optional) and any initialization params.

# Register the encoder
smp.encoders.encoders["my_awesome_encoder"] = {
    "encoder": MyEncoder, # encoder class here
    "pretrained_settings": {
        "imagenet": {
            "mean": [0.485, 0.456, 0.406],
            "std": [0.229, 0.224, 0.225],
            "url": "https://some-url.com/my-model-weights",
            "input_space": "RGB",
            "input_range": [0, 1],
        },
    },
    "params": {
        # init params for encoder if any
    },
}

# Use the encoder
model = smp.Unet(encoder_name="my_awesome_encoder")

To install the latest development version directly from the GitHub repository, use the following command.

$ pip install -U git+https://github.com/qubvel/segmentation_models.pytorch

To use the library, install segmentation-models-pytorch along with lightning and albumentations for training and augmentation support.

!pip install segmentation-models-pytorch lightning albumentations

To achieve better results and faster convergence, you should preprocess your input data using the same method used during the original weight pre-training. This is particularly relevant for 1, 2, or 3-channel images.

You can retrieve the appropriate preprocessing function using get_preprocessing_fn from segmentation_models_pytorch.encoders by providing the encoder name and the pre-training type.

from segmentation_models_pytorch.encoders import get_preprocessing_fn

preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')

You can install the library via PyPI or directly from the GitHub repository.

To install the stable version from PyPI:

pip install segmentation-models-pytorch

To install the latest version from GitHub:

pip install git+https://github.com/qubvel/segmentation_models.pytorch