Segmentation models is python library with Neural Networks for Image Segmentation based on Keras (Tensorflow) framework.
The main features of this library are:
- High level API (just two lines to create NN)
- 4 models architectures for binary and multi class segmentation (including legendary Unet)
- 25 available backbones for each architecture
- All backbones have pre-trained weights for faster and better convergence
Since the library is built on the Keras framework, created segmentaion model is just a Keras Model, which can be created as easy as:
from segmentation_models import Unet
model = 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 = Unet('resnet34', encoder_weights='imagenet')Change number of output classes in the model (choose your case):
# binary segmentation (this parameters are default when you call Unet('resnet34')
model = Unet('resnet34', classes=1, activation='sigmoid')# multiclass segmentation with non overlapping class masks (your classes + background)
model = Unet('resnet34', classes=3, activation='softmax')# multiclass segmentation with independent overlapping/non-overlapping class masks
model = Unet('resnet34', classes=3, activation='sigmoid')Change input shape of the model:
# if you set input channels not equal to 3, you have to set encoder_weights=None
# how to handle such case with encoder_weights='imagenet' described in docs
model = Unet('resnet34', input_shape=(None, None, 6), encoder_weights=None)from segmentation_models import Unet
from segmentation_models.backbones import get_preprocessing
from segmentation_models.losses import bce_jaccard_loss
from segmentation_models.metrics import iou_score
BACKBONE = 'resnet34'
preprocess_input = get_preprocessing(BACKBONE)
# load your data
x_train, y_train, x_val, y_val = load_data(...)
# preprocess input
x_train = preprocess_input(x_train)
x_val = preprocess_input(x_val)
# define model
model = Unet(BACKBONE, encoder_weights='imagenet')
model.compile('Adam', loss=bce_jaccard_loss, metrics=[iou_score])
# fit model
# if you use data generator use model.fit_generator(...) instead of model.fit(...)
# more about `fit_generator` here: https://keras.io/models/sequential/#fit_generator
model.fit(
x=x_train,
y=y_train,
batch_size=16,
epochs=100,
validation_data=(x_val, y_val),
)Same manimulations can be done with Linknet, PSPNet and FPN. For more detailed information about models API and use cases Read the Docs.
Models
| Unet | Linknet |
|---|---|
 |
 |
| PSPNet | FPN |
|---|---|
 |
 |
Backbones
| Type | Names |
|---|---|
| VGG | 'vgg16' 'vgg19' |
| ResNet | 'resnet18' 'resnet34' 'resnet50' 'resnet101' 'resnet152' |
| SE-ResNet | 'seresnet18' 'seresnet34' 'seresnet50' 'seresnet101' 'seresnet152' |
| ResNeXt | 'resnext50' 'resnext101' |
| SE-ResNeXt | 'seresnext50' 'seresnext101' |
| SENet154 | 'senet154' |
| DenseNet | 'densenet121' 'densenet169' 'densenet201' |
| Inception | 'inceptionv3' 'inceptionresnetv2' |
| MobileNet | 'mobilenet' 'mobilenetv2' |
| EfficientNet | 'efficientnetb0' 'efficientnetb1' 'efficientnetb2' 'efficientnetb3' |
All backbones have weights trained on 2012 ILSVRC ImageNet dataset (encoder_weights='imagenet').
Requirements
- Python 3.5+
- Keras >= 2.2.0
- Keras Application >= 1.0.7
- Image Classifiers == 0.2.0
- Tensorflow 1.9 (tested)
Pip package
$ pip install segmentation-modelsLatest version
$ pip install git+https://github.com/qubvel/segmentation_modelsLatest documentation is avaliable on Read the Docs
To see important changes between versions look at CHANGELOG.md
@misc{Yakubovskiy:2019,
Author = {Pavel Yakubovskiy},
Title = {Segmentation Models},
Year = {2019},
Publisher = {GitHub},
Journal = {GitHub repository},
Howpublished = {\url{https://github.com/qubvel/segmentation_models}}
}
Project is distributed under MIT Licence.