README.md

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. Models
   1. Architectires
   2. Encoders
   3. Pretrained weights
3. Installation
4. 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('renset18', pretrained='imagenet')

Model training example on CamVid dataset here.

Models

Architectires

• Unet
• Linknet
• FPN
• PSPNet

Encoders

Weights

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