Segformer B5 640x640 Ade 160k

🚀 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"

# Load pretrained model and preprocessing function
checkpoint = "smp-hub/segformer-b5-640x640-ade-160k"
model = smp.from_pretrained(checkpoint).eval().to(device)
preprocessing = A.Compose.from_pretrained(checkpoint)

# Load image
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)

# Preprocess image
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)  # HWC -> BCHW
input_tensor = input_tensor.to(device)

# Perform inference
with torch.no_grad():
    output_mask = model(input_tensor)

# Postprocess mask
mask = torch.nn.functional.interpolate(
    output_mask, size=(image.height, image.width), mode="bilinear", align_corners=False
)
mask = mask.argmax(1).cpu().numpy()  # argmax over predicted classes (channels dim)

✨ 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

• Library: You can find more information about the segmentation_models_pytorch library at https://github.com/qubvel/segmentation_models.pytorch.
• Docs: Detailed documentation can be found at https://smp.readthedocs.io/en/latest/.

📄 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.