Segmentation models is python library with Neural Networks for Image Segmentation based on PyTorch.

The main features of this library are:

• High level API (just two lines to create neural network)
• 4 models architectures for binary and multi class segmentation (including legendary Unet)
• 30 available encoders for each architecture
• All encoders have pre-trained weights for faster and better convergence

1. Quick start
2. Examples
3. Models
   1. Architectures
   2. Encoders
   3. Pretrained weights
4. Models API
5. Installation
6. License

Since the library is built on the PyTorch framework, created segmentation model is just a PyTorch nn.Module, which can be created as easy as:

import segmentation_models_pytorch as smp

model = smp.Unet()

Depending on the task, you can change the network architecture by choosing backbones with fewer or more parameters and use pretrainded weights to initialize it:

model = smp.Unet('resnet34', encoder_weights='imagenet')

Change number of output classes in the model:

model = smp.Unet('resnet34', classes=3, activation='softmax')

All models have pretrained encoders, so you have to prepare your data the same way as during weights pretraining:

from segmentation_models_pytorch.encoders import get_preprocessing_fn

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

• Training model for cars segmentation on CamVid dataset here.
• Training model with Catalyst (high-level framework for PyTorch) - here.

• Unet
• Linknet
• FPN
• PSPNet

• model.encoder - pretrained backbone to extract features of different spatial resolution
• model.decoder - segmentation head, depends on models architecture (Unet/Linknet/PSPNet/FPN)
• model.activation - output activation function, one of sigmoid, softmax
• model.forward(x) - sequentially pass x through model`s encoder and decoder (return logits!)
• model.predict(x) - inference method, switch model to .eval() mode, call .forward(x) and apply activation function with torch.no_grad()

PyPI version:

$ pip install segmentation-models-pytorch

Latest version from source:

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

Project is distributed under MIT License

$ docker build -f docker/Dockerfile.dev -t smp:dev .
$ docker run --rm smp:dev pytest -p no:cacheprovider