Upernet Convnext Small

🚀 UPerNet Model Card

This is a model card for the UPerNet model used in image segmentation. It provides details on loading the trained model, model initialization parameters, and the dataset.

🚀 Quick Start

Load trained model

You can quickly start using the trained model by following these steps. Click the button below to open the example in Google Colab.

1. Install requirements.

pip install -U segmentation_models_pytorch albumentations

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/upernet-convnext-small"
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 = mask.argmax(1).cpu().numpy()  # argmax over predicted classes (channels dim)

✨ Features

Model init parameters

The following are the initialization parameters for the model.

model_init_params = {
    "encoder_name": "tu-convnext_small.in12k_ft_in1k",
    "encoder_depth": 5,
    "encoder_weights": "imagenet",
    "decoder_channels": 512,
    "decoder_use_norm": "batchnorm",
    "in_channels": 3,
    "classes": 150,
    "activation": None,
    "aux_params": None
}

Dataset

The model uses the ADE20K dataset.

📚 Documentation

More Information

• Library: https://github.com/qubvel/segmentation_models.pytorch
• Docs: https://smp.readthedocs.io/en/latest/

This model has been pushed to the Hub using the PytorchModelHubMixin

📄 License

This project is licensed under the MIT license.