Repository files navigation

The main features of the library are:

• Super simple high-level API (just two lines to create a neural network)
• 12 encoder-decoder model architectures (Unet, Unet++, Segformer, DPT, ...)
• 800+pretrained convolution- and transform-based encoders, includingtimm support
• Popular metrics and losses for training routines (Dice, Jaccard, Tversky, ...)
• ONNX export and torch script/trace/compile friendly

withoutBG API https://withoutbg.com High-quality background removal API

VisitRead The Docs Project Page or read the following README to know more about Segmentation Models Pytorch (SMP for short) library

1. Quick start
2. Examples
3. Models and encoders
4. Models API
   1. Input channels
   2. Auxiliary classification output
   3. Depth
5. Installation
6. Competitions won with the library
7. Contributing
8. Citing
9. License

The segmentation model is just a PyTorchtorch.nn.Module, which can be created as easy as:

importsegmentation_models_pytorchassmpmodel=smp.Unet(encoder_name="resnet34",# choose encoder, e.g. mobilenet_v2 or efficientnet-b7encoder_weights="imagenet",# use `imagenet` pre-trained weights for encoder initializationin_channels=1,# model input channels (1 for gray-scale images, 3 for RGB, etc.)classes=3,# model output channels (number of classes in your dataset))

• seetable with available model architectures
• seetable with available encoders and their corresponding weights

All encoders have pretrained weights. Preparing your data the same way as during weights pre-training may give you better results (higher metric score and faster convergence). It isnot necessary in case you train the whole model, not only the decoder.

fromsegmentation_models_pytorch.encodersimportget_preprocessing_fnpreprocess_input=get_preprocessing_fn('resnet18',pretrained='imagenet')

Congratulations! You are done! Now you can train your model with your favorite framework!

Name	Link	Colab
Train pets binary segmentation on OxfordPets	Notebook
Train cars binary segmentation on CamVid	Notebook
Train multiclass segmentation on CamVid	Notebook
Train clothes binary segmentation by @ternaus	Repo
Load and inference pretrained Segformer	Notebook
Load and inference pretrained DPT	Notebook
Load and inference pretrained UPerNet	Notebook
Save and load models locally / to HuggingFace Hub	Notebook
Export trained model to ONNX	Notebook

The library provides a wide range ofpretrained encoders (also known as backbones) for segmentation models. Instead of using features from the final layer of a classification model, we extractintermediate features and feed them into the decoder for segmentation tasks.

All encoders come withpretrained weights, which help achievefaster and more stable convergence when training segmentation models.

Given the extensive selection of supported encoders, you can choose the best one for your specific use case, for example:

• Lightweight encoders for low-latency applications or real-time inference on edge devices (mobilenet/mobileone).
• High-capacity architectures for complex tasks involving a large number of segmented classes, providing superior accuracy (convnext/swin/mit).

By selecting the right encoder, you can balanceefficiency, performance, and model complexity to suit your project needs.

All encoders and corresponding pretrained weight are listed in the documentation:

• table with natively ported encoders
• table withtimm encoders supported

The input channels parameter allows you to create a model that can process a tensor with an arbitrary number of channels.If you use pretrained weights from ImageNet, the weights of the first convolution will be reused:

• For the 1-channel case, it would be a sum of the weights of the first convolution layer.
• Otherwise, channels would be populated with weights likenew_weight[:, i] = pretrained_weight[:, i % 3], and then scaled withnew_weight * 3 / new_in_channels.

model=smp.FPN('resnet34',in_channels=1)mask=model(torch.ones([1,1,64,64]))

All models supportaux_params parameters, which is default set toNone.Ifaux_params = None then classification auxiliary output is not created, elsemodel produce not onlymask, but alsolabel output with shapeNC.Classification head consists of GlobalPooling->Dropout(optional)->Linear->Activation(optional) layers, which can beconfigured byaux_params as follows:

aux_params=dict(pooling='avg',# one of 'avg', 'max'dropout=0.5,# dropout ratio, default is Noneactivation='sigmoid',# activation function, default is Noneclasses=4,# define number of output labels)model=smp.Unet('resnet34',classes=4,aux_params=aux_params)mask,label=model(x)

Depth parameter specify a number of downsampling operations in encoder, so you can makeyour model lighter if specify smallerdepth.

model=smp.Unet('resnet34',encoder_depth=4)

PyPI version:

$ pip install segmentation-models-pytorch

The latest version from GitHub:

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

Segmentation Models package is widely used in image segmentation competitions.Here you can find competitions, names of the winners and links to their solutions.

1. Install SMP in dev mode

make install_dev# Create .venv, install SMP in dev mode

2. Run tests and code checks

maketest# Run tests suite with pytestmake fixup# Ruff for formatting and lint checks

3. Update a table (in case you added an encoder)

make table# Generates a table with encoders and print to stdout

@misc{Iakubovskii:2019,  Author = {Pavel Iakubovskii},  Title = {Segmentation Models Pytorch},  Year = {2019},  Publisher = {GitHub},  Journal = {GitHub repository},  Howpublished = {\url{https://github.com/qubvel/segmentation_models.pytorch}}}

The project is primarily distributed underMIT License, while some files are subject to other licenses. Please refer toLICENSES and license statements in each file for careful check, especially for commercial use.