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