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
- Quick start
- Examples
- Models
- Models API
- Installation
- Competitions won with the library
- Contributing
- Citing
- License
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')
- 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) - here