Skip to main content

# Segmentation models
[![Build Status](https://travis-ci.com/qubvel/segmentation_models.pytorch.svg?branch=master)](https://travis-ci.com/qubvel/segmentation_models.pytorch) [![Generic badge](https://img.shields.io/badge/License-MIT-<COLOR>.svg)](https://shields.io/)

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](#start)
2. [Models](#models)
1. [Architectires](#architectires)
2. [Encoders](#encoders)
3. [Pretrained weights](#weights)
3. [Installation](#installation)
4. [License](#license)

### Quick start <a name="start"></a>
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:
```python
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:

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

Change number of output classes in the model:

```python
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:
```python
from segmentation_models_pytorch.encoders import get_preprocessing_fn

preprocess_input = get_preprocessing_fn('renset18', pretrained='imagenet')
```

**Model training example** on CamVid dataset [here](https://github.com/qubvel/segmentation_models.pytorch/blob/master/examples/cars%20segmentation%20(camvid).ipynb).

### Models <a name="models"></a>

#### Architectires <a name="architectires"></a>
- [Unet](https://arxiv.org/abs/1505.04597)
- [Linknet](https://arxiv.org/abs/1707.03718)
- [FPN](http://presentations.cocodataset.org/COCO17-Stuff-FAIR.pdf)
- [PSPNet](https://arxiv.org/abs/1612.01105)

#### Encoders <a name="encoders"></a>

| Type | Encoder names |
|------------|-----------------------------------------------------------------|
| VGG | vgg11, vgg13, vgg16, vgg19, vgg11bn, vgg13bn, vgg16bn, vgg19bn |
| DenseNet | densenet121, densenet169, densenet201, densenet161 |
| DPN | dpn68, dpn68b, dpn92, dpn98, dpn107, dpn131 |
| Inception | inceptionresnetv2 |
| ResNet | resnet18, resnet34, resnet50, resnet101, resnet152 |
| SE-ResNet | se_resnet50, se_resnet101, se_resnet152 |
| SE-ResNeXt | se_resnext50_32x4d, se_resnext101_32x4d |
| SENet | senet154 | |

#### Weights <a name="weights"></a>

| Weights name | Encoder names |
|--------------|-----------------------|
| imagenet+5k | dpn68b, dpn92, dpn107 |
| imagenet | * all other encoders |


### Installation <a name="installation"></a>
PyPI version is not released yet, to install package from source run following command:
```bash
$ pip install git+https://github.com/qubvel/segmentation_models.pytorch
````
### License <a name="license"></a>
Project is distributed under [MIT License](https://github.com/qubvel/segmentation_models.pytorch/blob/master/LICENSE)

### Run tests
```bash
$ docker build -f docker/Dockerfile.dev -t smp:dev .
$ docker run --rm smp:dev pytest -p no:cacheprovider
```


Release files for segmentation-models-pytorch 0.0.1

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for segmentation-models-pytorch 0.0.1
File Size Uploaded
segmentation_models_pytorch-0.0.1.tar.gz 16.8 kB Details

Built distribution (wheel)

Table of built distributions (wheels) for segmentation-models-pytorch 0.0.1
File Interpreter ABI Platform
segmentation_models_pytorch-0.0.1-py2.py3-none-any.whl Python 3, Python 2 none any Details

Total release size: 40.9 kB

Release files / segmentation_models_pytorch-0.0.1.tar.gz

Download URL segmentation_models_pytorch-0.0.1.tar.gz
Size 16.8 kB
Tags Source
SHA-256 checksum
How to use checksums
f4b2338f995565f9a0018ac869fc7240569d87a4a63dfb666ed59533c7405423
BLAKE2b-256 checksum
How to use checksums
525c60d03e239f3c49ab293b4c26b16262fcde77ca1eb95d8b794f75f9e0ecf0
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/1.13.0 pkginfo/1.5.0.1 requests/2.18.4 setuptools/39.2.0 requests-toolbelt/0.9.1 tqdm/4.19.9 CPython/3.6.5

Release files / segmentation_models_pytorch-0.0.1-py2.py3-none-any.whl

Download URL segmentation_models_pytorch-0.0.1-py2.py3-none-any.whl
Size 24.1 kB
Tags Python 2 Python 3
SHA-256 checksum
How to use checksums
694e7985d98e2a1a2caece322c5fa3c4b2b7fcc2c78cd00afb0372026a4a62cb
BLAKE2b-256 checksum
How to use checksums
537fa009f9d116ca46be5ce8be2e2de318c4da57b62e32fa0b11a938b2808bb8
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/1.13.0 pkginfo/1.5.0.1 requests/2.18.4 setuptools/39.2.0 requests-toolbelt/0.9.1 tqdm/4.19.9 CPython/3.6.5

Release history Release notifications | RSS feed

0.5.0

2 release files

0.4.0

2 release files

0.3.4

2 release files

0.3.3

2 release files

0.3.2

2 release files

0.3.1

2 release files

0.3.0

2 release files

0.2.1

2 release files

0.2.0

2 release files

0.1.3

2 release files

0.1.2

2 release files

0.1.1

2 release files

0.1.0

2 release files

0.0.3

2 release files

0.0.2

1 release file

This release

0.0.1 This release

2 release files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page