Registry / ai-ml / segmentation-models-pytorch

segmentation-models-pytorch

library0.5.0pypypi✓ verified 22d ago

Segmentation Models PyTorch (SMP) is a Python library offering a high-level API for various neural network architectures, pre-trained backbones, losses, and metrics for image semantic segmentation. It supports 12 encoder-decoder architectures and over 800 pre-trained convolutional and transformer-based encoders, leveraging `timm` for a vast selection. The library focuses on simplicity, fast convergence, and compatibility with PyTorch's `torch.jit.script`, `torch.compile`, and `torch.export` features. It is currently at version 0.5.0 and maintains an active release cadence with frequent updates and new model integrations.

ai-ml

pip install segmentation-models-pytorch

github ↗package ↗

INSTALL

IMPORT

SIG · SEGMENTATION-MODEL

segmentation-models-pytorch

ai-mlpythonv0.5.0

Install

75.2s avg

Import

15058ms

Disk

2182MB

Pass rate

9/ 10

Env Coverage9 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.0.2 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

✓ —

✓ 83.8s

py 3.11

✓ —

✓ 77.3s

py 3.12

✓ —

✓ 72s

py 3.13

✓ —

✓ 67.5s

py 3.9

✓ —

✕ timeout

2182MB installed

● package 2182MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

smp

✓ import segmentation_models_pytorch as smp

Unet

✓ model = smp.Unet(...)

✗ from segmentation_models_pytorch.models import Unet

Architectures are directly accessible from the top-level `smp` module for convenience.

get_preprocessing_fn

✓ from segmentation_models_pytorch.encoders import get_preprocessing_fn

This quickstart demonstrates how to instantiate a U-Net model with a pre-trained ResNet34 encoder, configure input and output channels, and set an activation function. It also shows how to obtain and apply the correct preprocessing function required for models with pre-trained backbones to ensure optimal performance.

import torch import segmentation_models_pytorch as smp from segmentation_models_pytorch.encoders import get_preprocessing_fn # 1. Create segmentation model model = smp.Unet( encoder_name="resnet34", # choose encoder backbone encoder_weights="imagenet", # use `imagenet` pre-trained weights in_channels=3, # model input channels (3 for RGB) classes=1, # model output channels (number of classes) activation='sigmoid' # activation function for binary segmentation ) # 2. Configure data preprocessing (important for pre-trained encoders) preprocess_input = get_preprocessing_fn('resnet34', pretrained='imagenet') # Example usage: # Dummy input image (batch_size=1, channels=3, height=256, width=256) image = torch.randn(1, 3, 256, 256) # Apply preprocessing (e.g., normalization) input_tensor = preprocess_input(image) # Forward pass model.eval() with torch.no_grad(): predicted_mask = model(input_tensor) print(f"Model output shape: {predicted_mask.shape}")

Debug

Known issues

breakingThe `UperNet` model architecture underwent significant changes in v0.5.0, making model weights trained with v0.4.0 incompatible with v0.5.0. Existing UperNet models will need to be re-trained or adapted.

fix

Re-train UperNet models with SMP v0.5.0 or ensure strict versioning. For fine-tuning, load the new `smp-hub` checkpoints.

affects: >=0.5.0

deprecatedEncoders from the `timm` library previously accessed with a `timm-` prefix (e.g., `timm-resnet34`) are deprecated in v0.5.0. The recommended way to use `timm` encoders is now with the `tu-` prefix (e.g., `tu-resnet34`).

fix

Update `encoder_name` arguments to use the `tu-` prefix for all `timm` encoders.

affects: >=0.5.0

deprecatedThe `smp.utils.losses` module was deprecated in v0.2.0. All loss functions have been moved to the `smp.losses` module.

fix

Update import paths and usage from `smp.utils.losses` to `smp.losses`.

affects: >=0.2.0

breakingThe minimum Python version requirement was increased from 3.6 to 3.7.

fix

Ensure your environment uses Python 3.7 or newer.

affects: >=0.3.2

gotchaTo ensure compatibility with `albumentations` versions >= 1.4.0, some internal function names that interact with `albumentations` may have changed, requiring updates if you directly extended or modified SMP's data processing pipelines.

fix

Review release notes for v0.3.4 and `albumentations` documentation for affected function names if custom integrations are breaking. Upgrade `albumentations` if issues persist.

affects: >=0.3.4

gotchaFor optimal performance, especially when using pre-trained encoders, it is crucial to apply the correct preprocessing steps (e.g., normalization, resizing) to your input data, matching how the encoder's weights were pre-trained. Use `smp.encoders.get_preprocessing_fn` for this.

fix

Always integrate `get_preprocessing_fn` into your data loading pipeline to preprocess images before feeding them to the model.

affects: All versions

gotchaSome models may require input image dimensions to be a power of 2, or they may handle incorrect sizes with specific interpolation/padding methods. Unexpected input sizes can lead to errors or degraded performance.

fix

Ensure input image height and width are compatible with the chosen model architecture and encoder, often powers of 2 (e.g., 256x256, 512x512). Refer to specific model documentation for exact requirements.

affects: All versions

Upgrade

Version history

0.5.0latest on PyPI · released Apr 17, 2025

Audit

Dependencies

torchrequiredCore PyTorch dependency for model definition and execution.

timmrequiredProvides a large collection of pre-trained image models used as encoders (backbones).

numpyrequiredFundamental package for numerical computing.

pillowrequiredImage processing library, often used in data loading.

torchvisionoptionalProvides datasets, models, and image transformations for computer vision.

albumentationsoptionalPopular library for image augmentation, frequently used with SMP for training.

huggingface-huboptionalFor saving, loading, and sharing models with Hugging Face Hub (added in v0.3.4).

Agent activity

13 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/qubvel-org/segmentation_models.pytorch ↗

packagehttps://pypi.org/project/segmentation-models-pytorch/ ↗