Segmentation Models - v0.5.0

New Models

DPT

The DPT model adapts the Vision Transformer (ViT) architecture for dense prediction tasks like semantic segmentation. It uses a ViT as a powerful backbone, processing image information with a global receptive field at each stage. The key innovation lies in its decoder, which reassembles token representations from various transformer stages into image-like feature maps at different resolutions. These are progressively combined using convolutional PSP and FPN blocks to produce full-resolution, high-detail predictions.

The model in smp can be used with a wide variety of transformer-based encoders

import segmentation_models_pytorch as smp

# initialize with your own pretrained encoder
model = smp.DPT("tu-mobilevitv2_175.cvnets_in1k", classes=2)

# load fully-pretrained on ADE20K 
model = smp.from_pretrained("smp-hub/dpt-large-ade20k")

# load the same checkpoint for finetuning
model = smp.from_pretrained("smp-hub/dpt-large-ade20k", classes=1, strict=False)

The full table of DPT's supported timm encoders can be found here.

• Adding DPT by @vedantdalimkar in #1079

Models export

A lot of work was done to add support for torch.jit.script, torch.compile (without graph breaks: fullgraph=True) and torch.export features in all encoders and models.

This provides several advantages:

• torch.jit.script: Enables serialization of models into a static graph format, enabling deployment in environments without a Python interpreter and allowing for graph-based optimizations.
• torch.compile (with fullgraph=True): Leverages Just-In-Time (JIT) compilation (e.g., via Triton or Inductor backends) to generate optimized kernels, reducing Python overhead and enabling significant performance improvements through techniques like operator fusion, especially on GPU hardware. fullgraph=True minimizes graph breaks, maximizing the scope of these optimizations.
• torch.export: Produces a standardized Ahead-Of-Time (AOT) graph representation, simplifying the process of exporting models to various inference backends and edge devices (e.g., through ExecuTorch) while preserving model dynamism where possible.

PRs:

• Fix torch compile, script, export by @qubvel in #1031
• Fix Efficientnet encoder for torchscript by @qubvel in #1037

Core

All encoders from third-party libraries such as efficientnet-pytorch and pretrainedmodels.pytorch are now vendored by SMP. This means we have copied and refactored the underlying code and moved all checkpoints to the smp-hub. As a result, you will have fewer additional dependencies when installing smp and get much faster weights downloads.

• Move encoders weights to HF-Hub by @qubvel in #1035
• Vendor pretrainedmodels by @adamjstewart in #1039
• Vendor efficientnet-pytorch by @adamjstewart in #1036

🚨🚨🚨 Breaking changes

1. UperNet model was significantly changed to reflect the original implementation and to bring pretrained checkpoints into SMP. Unfortunately, UperNet model weights trained with v0.4.0 will be not compatible with SMP v0.5.0.

   • Fix UperNet model and add pretrained checkpoints by @qubvel in #1124

2. While the high-level API for modeling should be backward compatible with v0.4.0, internal modules (such as encoders, decoders, blocks) might have changed initialization and forward interfaces.

3. timm- prefixed encoders are deprecated, tu- variants are now the recommended way to use encoders from the timm library. Most of the timm- encoders are internally switched to their tu- equivalent with state_dict re-mapping (backward-compatible), but this support will be dropped in upcoming versions.

Other changes

• Enable any resolution for Unet by @qubvel in #1029
• Update README.md by @qubvel in #1046
• Add binary segmentation example using cpu by @omidvarnia in #1057
• Load model with mismatched sizes by @qubvel in #1107
• Deprecate use_batchnorm in favor of generalized use_norm parameter by @GuillaumeErhard in #1095
• Extend usage of interpolation_mode to MAnet / UnetPlusPlus / FPN and align PAN by @GuillaumeErhard in #1108
• Fix cls token slicing for DPT by @qubvel in #1121
• add upsampling parameter #1106 by @DCalhas in #1123
• Fix #1125 by @Fede1995 in #1126

New Contributors

• @omidvarnia made their first contribution in #1057
• @GuillaumeErhard made their first contribution in #1095
• @kocabiyik made their first contribution in #1113
• @vedantdalimkar made their first contribution in #1079
• @DCalhas made their first contribution in #1123
• @Fede1995 made their first contribution in #1126

Full Changelog: v0.4.0...v0.5.0

Segmentation Models - v0.4.0

New models

Segformer

contributed by @brianhou0208

Originally, SegFormer is a transformer-based semantic segmentation model known for its simplicity and efficiency. It uses a lightweight hierarchical encoder to capture multi-scale features and a minimal decoder for fast inference.

With segmentation-models-pytorch you can utilize the model with a native Mix Vision Transformer encoder as long as with 800+ other encoders supported by the library. Original weights are also supported and can be loaded as follows:

import segmentation_models_pytorch as smp

model = smp.from_pretrained("smp-hub/segformer-b5-640x640-ade-160k")

or with any other encoder:

import segmentation_models_pytorch as smp

model = smp.Segformer("resnet34")

See more checkpoints on the HF Hub.

UperNet

contributed by @brianhou0208

UPerNet (Unified Perceptual Parsing Network) is a versatile semantic segmentation model designed to handle diverse scene parsing tasks. It combines a Feature Pyramid Network (FPN) with a Pyramid Pooling Module (PPM) to effectively capture multi-scale context.

import segmentation_models_pytorch as smp

model = smp.UPerNet("resnet34")

New Encoders

Thanks to @brianhou0208 contribution 800+ timm encoders are now supported in segmentation_models.pytorch. New modern encoders like convnext, efficientvit, efficientformerv2, hiera, mambaout and more can be used as easy as:

import segmentation_models_pytorch as smp

model = smp.create_model("upernet", encoder_name="tu-mambaout_small")
# or
model = smp.UPerNet("tu-mambaout_small")

New examples

• Added example for multi-class segmentation by @TimbusCalin
• Added example for onnx export by @qubvel

Other changes

• Project migrated to pyproject.toml by @adamjstewart
• Better dependency managing and testing (minimal and latest dependencies, linux/windows/mac platforms) by @adamjstewart
• Better type annotations
• Tests are refactored for faster CI and local testing by @qubvel

All changes

• Updating the tutorial file by @ytzfhqs in #907
• Example on how to save and load model along with Albumentations preprocessing by @qubvel in #914
• Add open-in-colab badge for all example notebooks by @qubvel in #915
• Switch to pyproject.toml by @adamjstewart in #917
• Remove dep on mock by @adamjstewart in #919
• [feat] Adding camvid segmentation multiclass as an example by @TimbusCalin in #922
• Ruff: format Jupyter notebooks too by @adamjstewart in #923
• Remove docker files by @adamjstewart in #925
• Test minimum and maximum supported dependencies by @adamjstewart in #918
• Test on Linux/macOS/Windows for all supported Python versions by @adamjstewart in #930
• Modify Jaccard, Dice and Tversky losses by @zifuwanggg in #927
• [feat] Adding UPerNet by @brianhou0208 in #926
• Fix dims=None in loss by @qubvel in #937
• Test PR docs build and update models.rst by @qubvel in #943
• Update test_models.py by @brianhou0208 in #940
• Fix UPerNet decoder typo by @brianhou0208 in #945
• Fix Metric typo by @brianhou0208 in #966
• Expose timm constructor arguments by @DimitrisMantas in #960
• fix(examples): correct Colab links by @EDM115 in #965
• Update DeepLab models by @DimitrisMantas in #959
• [feat] Adding SegFormer by @brianhou0208 in #944
• Update MixVisionTransformer by @brianhou0208 in #975
• silance "is" with 'str' literal syntax warning from pretrainedmodels in python >= 3.12 by @YoniChechik in #987
• Fix DeepLabV3Plus encoder depth by @munehiro-k in #986
• Fix style by @qubvel in #989
• Add onnx tutorial by @qubvel in #990
• Fix Segformer decoder performance by @brianhou0208 in #998
• Add description for non-MIT licensed codes by @junkoda in #1000
• Fix encoder depth & output stride on DeeplabV3 & DeeplabV3+ by @brianhou0208 in #991
• Update PAN Decoder support encoder depth by @brianhou0208 in #999
• Update timm universal (support transformer-style model) by @brianhou0208 in #1004
• Refactor tests by @qubvel in #1011
• Dependencies: packaging required for testing by @adamjstewart in #1013
• chore (ci): adopt astral-sh actions by @johnnv1 in #1014
• chore (segformer): move decoder converter scripts by @johnnv1 in #1017

New Contributors

• @adamjstewart made their first contribution in #917
• @TimbusCalin made their first contribution in #922
• @zifuwanggg made their first contribution in #927
• @brianhou0208 made their first contribution in #926
• @DimitrisMantas made their first contribution in #960
• @EDM115 made their first contribution in #965
• @YoniChechik made their first contribution in #987
• @munehiro-k made their first contribution in #986
• @junkoda made their first contribution in #1000
• @johnnv1 made their first contribution in #1014

Full Changelog: v0.3.4...v0.4.0

Segmentation Models - v0.3.4

Updates

• 🤗 Hugging Face integration: you can save, load, and share models with HF Hub, see example notebook.

Full log

• To support albumentations >= 1.4.0 some functions need to be renamed by @CallShaul in #870
• Updated false positve and false negative rate functions in functional.py by @vermavinay982 in #855
• Add HF hub mixin by @qubvel in #876
• use precommit for code linting by @Borda in #829
• Add Ruff for formatting and linting by @qubvel in #877
• Add docs config by @qubvel in #878
• Update docs by @qubvel in #879
• Add create_model to docs by @qubvel in #883
• Update ruff to version 0.5.2 and workflows update by @Smartappli in #892
• Fix hub_mixin.py pop error by @ytzfhqs in #909
• Update HF mixin by @qubvel in #910

New Contributors

• @CallShaul made their first contribution in #870
• @vermavinay982 made their first contribution in #855
• @Borda made their first contribution in #829
• @Smartappli made their first contribution in #892
• @ytzfhqs made their first contribution in #909

Full Changelog: v0.3.3...v0.3.4

Segmentation Models - v0.3.3

Updates

• Pytorch image models (timm) version upgrade to 0.9.2

Segmentation Models - v0.3.2

Updates

• Added Apple's Mobile One encoder from repo (use encoder_name="mobileone_s{0..4}").
• Pytorch image models (timm) version upgrade to 0.6.12 (500+ encoders available)
• Minor typo fixes and docs updates

Breaking changes

• Minimum Python version 3.6 -> 3.7

Thanks @VadimLevin, @kevinpl07, @Abd-elr4hman

Segmentation Models - v0.3.1

Updates

• Added Mix Vision Transformer encoder from SegFormer [official code] [paper]. Use argument encoder_name="mit_b0" (or mit_b1..b5) to create a model.
• Minor typo fixes and docs updates

Segmentation Models - v0.3.0

Updates

• Added smp.metrics module with different metrics based on confusion matrix, see docs
• Added new notebook with training example using pytorch-lightning
• Improved handling of incorrect input image size error (checking image size is 2^n)
• Codebase refactoring and style checks (black, flake8)
• Minor typo fixes and bug fixes

Breaking changes

• utils module is going to be deprecated, if you still need it import it manually from segmentation_models_pytorch import utils

Thanks a lot for all contributors!

Segmentation Models - v0.2.1

Updates

• Universal timm encoder. 400+ pretrained encoders from timm available with tu- prefix. List of available encoders here.
• Minor fixes and improvements.

Segmentation Models - v0.2.0

Updates

• New architecture: MANet (#310)
• New encoders from timm: mobilenetv3 (#355) and gernet (#344)
• New loss functions in smp.losses module (smp.utils.losses would be deprecated in future versions)
• New pretrained weight initialization for first convolution if in_channels > 3
• Updated timm version (0.4.12)
• Bug fixes and docs improvement

Thanks to @azkalot1 @JulienMaille @originlake @Kupchanski @loopdigga96 @zurk @nmerty @ludics @Vozf @markson14 and others!

Segmentation Models - v0.1.3

Updates

• New architecture Unet++ (#279)
• New encoders RegNet, ResNest, SK-Net, Res2Net (#286)
• Updated timm version (0.3.2)
• Improved docstrings and typehints for models
• Project documentation on https://smp.readthedocs.io

Thanks to @azkalot1 for the new encoders and architecture!