Segformer B0 512x512 Ade 160k

🚀 Segformer Model Card

This is a model card for the Segformer model, which is used for image segmentation tasks.

🚀 Quick Start

You can load the trained Segformer model and run inference following these steps. Click the button below to open the example in Google Colab:

Installation

First, install the required libraries:

pip install -U segmentation_models_pytorch albumentations

Inference

Then, run the following Python code for 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-b0-512x512-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

The Segformer model has the following features:

• It can be used for semantic segmentation tasks.
• It supports pre - trained models for easy inference.

📦 Installation

To use the Segformer model, you need to install the necessary libraries:

pip install -U segmentation_models_pytorch albumentations

💻 Usage Examples

Basic Usage

The above code in the "Quick Start" section shows the basic usage of loading a pre - trained Segformer model and performing inference on an image.

Advanced Usage

You can customize the model initialization parameters according to your needs. Here is an example of model initialization parameters:

model_init_params = {
    "encoder_name": "mit_b0",
    "encoder_depth": 5,
    "encoder_weights": None,
    "decoder_segmentation_channels": 256,
    "in_channels": 3,
    "classes": 150,
    "activation": None,
    "aux_params": None
}

📚 Documentation

• Library: segmentation - models.pytorch
• Docs: smp Documentation

📄 License

The license information for this model can be found at License.

Dataset

The model is trained on the ADE20K dataset.

This model has been pushed to the Hub using the PytorchModelHubMixin