transformers - 💡(How to fix) Fix Discussion: optional RankSEG-style decoding for Transformers semantic segmentation post-processing

Hi Transformers maintainers,

I wanted to share a small downstream experiment around RankSEG-style decoding for semantic segmentation. The short version is: if a Transformers processor can expose resized semantic class probabilities before the final argmax, then users can try metric-aware post-processing methods such as RankSEG without changing the model, checkpoint, or preprocessing pipeline.

This is related to https://github.com/huggingface/transformers/issues/37715, where the discussion is about making the final argmax optional and allowing users to access resized class probability maps. I do not want to assume that RankSEG itself belongs in Transformers, but I think it is a useful concrete example of why probability-level semantic segmentation outputs can matter.

What I Tried

RankSEG is a training-free segmentation decoding method. It takes per-class probability maps and returns a hard segmentation mask optimized for an overlap-style metric such as Dice or IoU. The relevant papers are RankSEG, JMLR 2023 and RankSEG-RMA, NeurIPS 2025. There is also a RankSEG repository, a PyPI package, and a Transformers integration tutorial.

The experiment used the usual Transformers inference path first:

inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)

Then I compared three post-processing choices using the same outputs:

# 1. Baseline: standard SegFormer / Transformers argmax-style decoding
upsampled_logits = torch.nn.functional.interpolate(
    outputs.logits,
    size=target_size,
    mode="bilinear",
    align_corners=False,
)
baseline = upsampled_logits.argmax(dim=1)[0]

# 2. RankSEG optimized for Dice
rankseg_dice = rankseg_transformers.postprocess(
    outputs,
    model=model,
    target_sizes=target_sizes,
    rankseg_kwargs={"metric": "dice", "solver": "RMA", "output_mode": "multiclass"},
)

# 3. RankSEG optimized for IoU
rankseg_iou = rankseg_transformers.postprocess(
    outputs,
    model=model,
    target_sizes=target_sizes,
    rankseg_kwargs={"metric": "iou", "solver": "RMA", "output_mode": "multiclass"},
)

The helper above is already implemented outside Transformers in RankSEG's current compatibility layer: documentation, source code, example script, and notebook. The same notebook can be opened in Colab.

Small Cityscapes Check

I used tanganke/cityscapes only as a lightweight local check because it has a convenient segmentation_19 ground-truth column. This is not an official Cityscapes benchmark. It is a small smoke test over the first 100 validation images, using samplewise macro Dice and IoU over non-empty classes.

The result is modest, but it is consistent with the intended use case: the model is unchanged, and only the final decoding step changes.

Visual Examples

Each image below uses the same layout: baseline argmax on the top left, RankSEG optimized for Dice on the top right, ground-truth overlay on the bottom left, and RankSEG optimized for IoU on the bottom right.

Why This Relates to Transformers Post-Processing

For simple semantic segmentation heads, restoring probabilities may look like resizing logits and applying softmax. For other model families, the post-processing path can involve class-query logits, mask logits, null classes, model-specific resizing conventions, or processor-owned logic. That is why a probability-returning option inside the existing Transformers post-processing API would be useful: the model-family-specific restoration would stay in the official processor path, while downstream methods could consume the restored probabilities.

Hard segmentation maps could remain the default behavior. The probability path would simply make the intermediate semantic distribution available for downstream decoding, calibration, uncertainty estimation, or metric-aware post-processing such as RankSEG.

Closing

I understand that adding or changing post-processing APIs has maintenance costs, especially in a library used across many model families. I am not asking maintainers to adopt RankSEG directly. I mainly wanted to share a concrete downstream use case showing why resized semantic probability maps could be useful to users who want to experiment beyond argmax.

I would also like to thank @statmlben and @ZixunWang, the RankSEG maintainers and authors of the recent RankSEG-RMA work, for developing and maintaining the RankSEG project that made this small Transformers experiment possible.

If maintainers think this direction is worth exploring, I would be happy to adapt the experiment to a preferred model family, test against a proposed API, or help write documentation/examples in the style that fits Transformers.