Skip to main content

Image segmentation models with pre-trained backbones with Keras.

Project description

PyPI version

Segmentation models Zoo

Segmentation models with pretrained backbones

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

$ git clone https://github.com/qubvel/segmentation_models.git
$ cd segmentation_models
$ git submodule update --init --recursive

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 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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

segmentation_models-0.1.2.tar.gz (28.2 kB view hashes)

Uploaded source

Built Distribution

segmentation_models-0.1.2-py2.py3-none-any.whl (45.6 kB view hashes)

Uploaded py2 py3

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor NVIDIA NVIDIA PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page