qubvel / segmentation_models

.. raw:: html Python library with Neural Networks for Image Segmentation based on Keras and TensorFlow . **The main features** of this library are: • High level API (just two lines of code to create model for segmentation) • **4** models architectures for binary and multi-class image segmentation (including legendary **Unet**) • **25** available backbones for each architecture • All backbones have **pre-trained** weights for faster and better convergence • Helpful segmentation losses (Jaccard, Dice, Focal) and metrics (IoU, F-score) **Important note** Some models of version 1.* are not compatible with previously trained models, if you have such models and want to load them - roll back with: $ pip install -U segmentation-models==0.2.1 Table of Contents ~~~~~~~~~~~~~~~~~ • _ • _ • _ • _ • _ • _ • _ • _ • _ Quick start ~~~~~~~~~~~ Library is build to work together with Keras and TensorFlow Keras frameworks .. code:: python import segmentation_models as sm # Segmentation Models: using framework. By default it tries to import keras , if it is not installed, it will try to start with tensorflow.keras framework. There are several ways to choose framework: • Provide environment variable SM_FRAMEWORK=keras / SM_FRAMEWORK=tf.keras before import segmentation_models • Change framework sm.set_framework('keras') / sm.set_framework('tf.keras') You can also specify what kind of image_data_format to use, segmentation-models works with both: channels_last and channels_first . This can be useful for further model conversion to Nvidia TensorRT format or optimizing model for cpu/gpu computations. .. code:: python import keras # or from tensorflow import keras keras.backend.set_image_data_format('channels_last') # or keras.backend.set_image_data_format('channels_first') Created segmentation model is just an instance of Keras Model, which can be build as easy as: .. code:: python model = sm.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: .. code:: python model = sm.Unet('resnet34', encoder_weights='imagenet') Change number of output classes in the model (choose your case): .. code:: python # binary segmentation (this parameters are default when you call Unet('resnet34') model = sm.Unet('resnet34', classes=1, activation='sigmoid') .. code:: python # multiclass segmentation with non overlapping class masks (your classes + background) model = sm.Unet('resnet34', classes=3, activation='softmax') .. code:: python # multiclass segmentation with independent overlapping/non-overlapping class masks model = sm.Unet('resnet34', classes=3, activation='sigmoid') Change input shape of the model: .. code:: python # 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) Simple training pipeline ~~~~~~~~~~~~~~~~~~~~~~~~ .. code:: python import segmentation_models as sm BACKBONE = 'resnet34' preprocess_input = sm.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 = sm.Unet(BACKBONE, encoder_weights='imagenet') model.compile( 'Adam', loss=sm.losses.bce_jaccard_loss, metrics=[sm.metrics.iou_score], ) # fit model # if you use data generator use model.fit_generator(...) instead of model.fit(...) # more about 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 manipulations can be done with Linknet , PSPNet and FPN . For more detailed information about models API and use cases __. Examples ~~~~~~~~ Models training examples: • [Jupyter Notebook] Binary segmentation ( ) on CamVid dataset __. • [Jupyter Notebook] Multi-class segmentation ( , ) on CamVid dataset __. Models and Backbones ~~~~~~~~~~~~~~~~~~~~ **Models** • __ • __ • __ • __ ============= ============== Unet Linknet ============= ============== |unet_image| |linknet_image| ============= ============== ============= ============== PSPNet FPN ============= ============== |psp_image| |fpn_image| ============= ============== .. _Unet: https://github.com/qubvel/segmentation_models/blob/readme/LICENSE .. _Linknet: https://arxiv.org/abs/1707.03718 .. _PSPNet: https://arxiv.org/abs/1612.01105 .. _FPN: http://presentations.cocodataset.org/COCO17-Stuff-FAIR.pdf .. |unet_image| image:: https://github.com/qubvel/segmentation_models/blob/master/images/unet.png .. |linknet_image| image:: https://github.com/qubvel/segmentation_models/blob/master/images/linknet.png .. |psp_image| image:: https://github.com/qubvel/segmentation_models/blob/master/images/pspnet.png .. |fpn_image| image:: https://github.com/qubvel/segmentation_models/blob/master/images/fpn.png **Backbones** .. table:: ============= ===== 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' 'efficientnetb4' 'efficientnetb5' efficientnetb6' efficientnetb7' ============= ===== .. epigraph:: All backbones have weights trained on 2012 ILSVRC ImageNet dataset ( encoder_weights='imagenet' ). Installation ~~~~~~~~~~~~ **Requirements** 1) python 3 2) keras >= 2.2.0 or tensorflow >= 1.13 3) keras-applications >= 1.0.7, <=1.0.8 4) image-classifiers ==…