Segmentation models

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

Table of content

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

Quick start

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')

Examples

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

Models

Architectures

• Unet
• Linknet
• FPN
• PSPNet

Encoders

Weights

Models API

• 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()

Installation

PyPI version:

$ pip install segmentation-models-pytorch

Latest version from source:

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

License

Project is distributed under MIT License

Run tests

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