segmentation_models.pytorch

Python library with Neural Networks for Image Semantic Segmentation based on PyTorch.

The main features of the library are:

• Super simple high-level API (just two lines to create a neural network)
• 12 encoder-decoder model architectures (Unet, Unet++, Segformer, DPT, ...)
• 800+ pretrained convolution- and transform-based encoders, including timm support
• Popular metrics and losses for training routines (Dice, Jaccard, Tversky, ...)
• ONNX export and torch script/trace/compile friendly

🤝 Sponsor: withoutBG

withoutBG is a high-quality background removal tool. They built their open-source image matting and refiner models using smp.Unet and are proudly sponsoring this project.

📚 Project Documentation 📚

Visit Read The Docs Project Page or read the following README to know more about Segmentation Models Pytorch (SMP for short) library

📋 Table of content

1. Quick start
2. Examples
3. Models and encoders
4. Models API
   1. Input channels
   2. Auxiliary classification output
   3. Depth
5. Installation
6. Competitions won with the library
7. Contributing
8. Citing
9. License

⏳ Quick start

1. Create your first Segmentation model with SMP

The segmentation model is just a PyTorch torch.nn.Module, which can be created as easy as:

import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=3,                      # model output channels (number of classes in your dataset)
)

• see table with available model architectures
• see table with available encoders and their corresponding weights

2. Configure data preprocessing

All encoders have pretrained weights. Preparing your data the same way as during weights pre-training may give you better results (higher metric score and faster convergence). It is not necessary in case you train the whole model, not only the decoder.

from segmentation_models_pytorch.encoders import get_preprocessing_fn

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

Congratulations! You are done! Now you can train your model with your favorite framework!

以上为 README 的开头部分。其余内容请阅读原文。 查看完整 README ↗