/segmentation_models

Segmentation models with pretrained backbones. Keras.

Primary LanguagePythonMIT LicenseMIT

PyPI version Documentation Status

Segmentation models Zoo

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)
  • 15 available backbones for each architecture
  • All backbones have pre-trained weights for faster and better convergence

Latest documentation is avaliable on Read the Docs

Avaliable models:

Avaliable backbones:

Backbone model Name Weights
VGG16 vgg16 imagenet
VGG19 vgg19 imagenet
ResNet18 resnet18 imagenet
ResNet34 resnet34 imagenet
ResNet50 resnet50 imagenet
imagenet11k-places365ch
ResNet101 resnet101 imagenet
ResNet152 resnet152 imagenet
imagenet11k
ResNeXt50 resnext50 imagenet
ResNeXt101 resnext101 imagenet
DenseNet121 densenet121 imagenet
DenseNet169 densenet169 imagenet
DenseNet201 densenet201 imagenet
Inception V3 inceptionv3 imagenet
Inception ResNet V2 inceptionresnetv2 imagenet

Requirements

  1. Python 3.6 or higher
  2. Keras >=2.1.0
  3. Tensorflow >= 1.4

Installation

Installing via pip

$ pip install segmentation-models

Using latest version in your project

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

Code examples

Train Unet model:

from segmentation_models import Unet
from segmentation_models.backbones import get_preprocessing

# prepare data
x, y = ...

preprocessing_fn = get_preprocessing('resnet34')
x = preprocessing_fn(x)

# prepare model
model = Unet(backbone_name='resnet34', encoder_weights='imagenet')
model.compile('Adam', 'binary_crossentropy', ['binary_accuracy'])

# train model
model.fit(x, y)

Train FPN model:

from segmentation_models import FPN

model = FPN(backbone_name='resnet34', encoder_weights='imagenet')

Useful trick

Freeze encoder weights for fine-tuning during first epochs of training:

from segmentation_models import FPN
from segmentation_models.utils import set_trainable

model = FPN(backbone_name='resnet34', encoder_weights='imagenet', freeze_encoder=True)
model.compile('Adam', 'binary_crossentropy', ['binary_accuracy'])

# pretrain model decoder
model.fit(x, y, epochs=2)

# release all layers for training
set_trainable(model) # set all layers trainable and recompile model

# continue training
model.fit(x, y, epochs=100)

TODO

  • Update Unet API
  • Update FPN API
  • Add Linknet models
  • Add PSP models
  • Add DPN backbones

Change Log

Version 0.1.2

Areas of improvement
  • Added PSPModel
  • Prepocessing functions for all backbones:
from segmentation_models.backbones import get_preprocessing

preprocessing_fn = get_preprocessing('resnet34')
X = preprocessing_fn(x)
API changes
  • Default param use_batchnorm=True for all decoders
  • FPN model Upsample2D layer renamed to ResizeImage

Version 0.1.1

  • Added Linknet model
  • Keras 2.2+ compatibility (fixed import of _obtain_input_shape)
  • Small code improvements and bug fixes

Version 0.1.0

  • Unet and FPN models