python图像分割模型_Segmentation models 是一个基于PyTorch的图像分割神经网络

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)

5 models architectures for binary and multi class segmentation (including legendary Unet)

46 available encoders for each architecture

All encoders have pre-trained weights for faster and better convergence

Table of content

Quick start

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:

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:

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

Change number of output classes in the model:

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:

from segmentation_models_pytorch.encoders import get_preprocessing_fn

preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')

Examples

Training model for cars segmentation on CamVid dataset here.

Training SMP model with Catalyst (high-level framework for PyTorch), Ttach (TTA library for PyTorch) and Albumentations (fast image augmentation library) -