Segmentation Models PyTorch

Segmentation Models PyTorch is a Python library of neural networks for image semantic segmentation, built on PyTorch, with a high-level API simple enough to create a segmentation model in two lines of code. It provides 12 encoder-decoder architectures, including Unet, Unet++, Segformer and DPT, and more than 800 pretrained convolutional and transformer-based encoders, including support for the timm library.

It includes commonly used segmentation metrics and losses, such as Dice, Jaccard and Tversky loss, for training a model, along with ONNX export and compatibility with TorchScript tracing and compilation. Encoders ship with pretrained weights and a matching preprocessing function, so a user's input data can be normalized the same way the encoder's own weights were originally trained, which the documentation notes can improve accuracy and convergence speed when only the decoder is being trained.

Its documentation covers a quick start, further examples, the available models and encoders, and installation instructions.