đ Segformer Model Card
This model card provides information about the Segformer model for image segmentation, including how to load the trained model, its initialization parameters, metrics, and the dataset used.
đ Quick Start
You can quickly start using the Segformer model by following these steps. Click the button below to open the example in Google Colab:

Step 1: Install requirements
pip install -U segmentation_models_pytorch albumentations
Step 2: Run inference
import torch
import requests
import numpy as np
import albumentations as A
import segmentation_models_pytorch as smp
from PIL import Image
device = "cuda" if torch.cuda.is_available() else "cpu"
checkpoint = "smp-hub/segformer-b5-640x640-ade-160k"
model = smp.from_pretrained(checkpoint).eval().to(device)
preprocessing = A.Compose.from_pretrained(checkpoint)
url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
image = Image.open(requests.get(url, stream=True).raw)
np_image = np.array(image)
normalized_image = preprocessing(image=np_image)["image"]
input_tensor = torch.as_tensor(normalized_image)
input_tensor = input_tensor.permute(2, 0, 1).unsqueeze(0)
input_tensor = input_tensor.to(device)
with torch.no_grad():
output_mask = model(input_tensor)
mask = torch.nn.functional.interpolate(
output_mask, size=(image.height, image.width), mode="bilinear", align_corners=False
)
mask = mask.argmax(1).cpu().numpy()
⨠Features
This Segformer model is designed for image segmentation tasks. It has the following features:
- Based on the
segmentation_models_pytorch
library, which provides a convenient way to use pre - trained models.
- Can be easily integrated into PyTorch projects.
đĻ Installation
The installation of the necessary libraries can be done using the following command:
pip install -U segmentation_models_pytorch albumentations
đģ Usage Examples
Basic Usage
The above code in the "Quick Start" section demonstrates the basic usage of the Segformer model for image segmentation inference.
Advanced Usage
You can customize the model by adjusting the initialization parameters. Here is an example of the model initialization parameters:
model_init_params = {
"encoder_name": "mit_b5",
"encoder_depth": 5,
"encoder_weights": None,
"decoder_segmentation_channels": 768,
"in_channels": 3,
"classes": 150,
"activation": None,
"aux_params": None
}
đ Documentation
đ License
The license information for this model can be found at https://github.com/NVlabs/SegFormer/blob/master/LICENSE.
Dataset
The model is trained on the ADE20K dataset.
This model has been pushed to the Hub using the PytorchModelHubMixin.