Discussion: optional RankSEG-style decoding for Transformers semantic segmentation post-processing
Author: Leev1s
Pull request: no

[![RankSEG](https://img.shields.io/badge/RankSEG-GitHub-blue?logo=github)](https://github.com/rankseg/rankseg) [![PyPI](https://badge.fury.io/py/rankseg.svg)](https://pypi.org/project/rankseg/) [![Docs](https://readthedocs.org/projects/rankseg/badge/?version=latest)](https://rankseg.readthedocs.io/en/latest/) [![Transformers docs](https://img.shields.io/badge/docs-Transformers%20integration-brightgreen)](https://rankseg.readthedocs.io/en/latest/integrations_transformers.html) [![Notebook](https://img.shields.io/badge/notebook-Transformers-orange)](https://github.com/rankseg/rankseg/blob/main/notebooks/rankseg_with_transformers.ipynb) [![JMLR 2023](https://img.shields.io/badge/JMLR-2023-black)](https://www.jmlr.org/papers/v24/22-0712.html) [![NeurIPS 2025](https://img.shields.io/badge/NeurIPS-2025-black)](https://openreview.net/forum?id=4tRMm1JJhw)

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](https://www.jmlr.org/papers/v24/22-0712.html) and [RankSEG-RMA, NeurIPS 2025](https://openreview.net/forum?id=4tRMm1JJhw). There is also a [RankSEG repository](https://github.com/rankseg/rankseg), a [PyPI package](https://pypi.org/project/rankseg/), and a [Transformers integration tutorial](https://rankseg.readthedocs.io/en/latest/integrations_transformers.html).

The experiment used the usual Transformers inference path first:

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

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

```python
# 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](https://rankseg.readthedocs.io/en/latest/integrations_transformers.html), [source code](https://github.com/rankseg/rankseg/blob/main/rankseg/integration/transformers.py), [example script](https://github.com/rankseg/rankseg/blob/main/examples/transformers_rankseg.py), and [notebook](https://github.com/rankseg/rankseg/blob/main/notebooks/rankseg_with_transformers.ipynb). The same notebook can be opened in [Colab](https://colab.research.google.com/github/rankseg/rankseg/blob/main/notebooks/rankseg_with_transformers.ipynb).

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

| Model | Method | Mean Dice | Dice delta | Mean IoU | IoU delta |
| --- | --- | ---: | ---: | ---: | ---: |
| `nvidia/segformer-b0-finetuned-cityscapes-512-1024` | Transformers argmax | 0.4608 | - | 0.3898 | - |
| `nvidia/segformer-b0-finetuned-cityscapes-512-1024` | RankSEG, `metric="dice"` | 0.4810 | +0.0202 | 0.4045 | +0.0147 |
| `nvidia/segformer-b0-finetuned-cityscapes-512-1024` | RankSEG, `metric="iou"` | 0.4813 | +0.0205 | 0.4051 | +0.0153 |
| `nvidia/segformer-b1-finetuned-cityscapes-1024-1024` | Transformers argmax | 0.4743 | - | 0.4015 | - |
| `nvidia/segformer-b1-finetuned-cityscapes-1024-1024` | RankSEG, `metric="dice"` | 0.4903 | +0.0160 | 0.4128 | +0.0113 |
| `nvidia/segformer-b1-finetuned-cityscapes-1024-1024` | RankSEG, `metric="iou"` | 0.4907 | +0.0164 | 0.4134 | +0.0118 |

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.

<table>
  <tr>
    <td width="50%" align="center">
      <a href="https://files.seeusercontent.com/2026/05/21/7wXl/rank_01_sample_0053_ddice_0092_d.png">
        <img src="https://files.seeusercontent.com/2026/05/21/7wXl/rank_01_sample_0053_ddice_0092_d.png" alt="SegFormer-B0 Cityscapes example 1" width="100%">
      </a>
      <br>
      <sub>SegFormer-B0, example 1</sub>
    </td>
    <td width="50%" align="center">
      <a href="https://files.seeusercontent.com/2026/05/21/aSr0/rank_02_sample_0029_ddice_0067_d.png">
        <img src="https://files.seeusercontent.com/2026/05/21/aSr0/rank_02_sample_0029_ddice_0067_d.png" alt="SegFormer-B0 Cityscapes example 2" width="100%">
      </a>
      <br>
      <sub>SegFormer-B0, example 2</sub>
    </td>
  </tr>
  <tr>
    <td width="50%" align="center">
      <a href="https://files.seeusercontent.com/2026/05/21/s6kI/rank_01_sample_0072_ddice_0084_d.png">
        <img src="https://files.seeusercontent.com/2026/05/21/s6kI/rank_01_sample_0072_ddice_0084_d.png" alt="SegFormer-B1 Cityscapes example 1" width="100%">
      </a>
      <br>
      <sub>SegFormer-B1, example 1</sub>
    </td>
    <td width="50%" align="center">
      <a href="https://files.seeusercontent.com/2026/05/21/Iyn9/rank_02_sample_0003_ddice_0081_d.png">
        <img src="https://files.seeusercontent.com/2026/05/21/Iyn9/rank_02_sample_0003_ddice_0081_d.png" alt="SegFormer-B1 Cityscapes example 2" width="100%">
      </a>
      <br>
      <sub>SegFormer-B1, example 2</sub>
    </td>
  </tr>
</table>

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


Comments:
Yeah, the lack of exposed probability maps is surprising! I'll try to push this internally

Thanks a lot @Rocketknight1, I really appreciate it!

Started a draft here: https://github.com/huggingface/transformers/pull/46163

Thanks @guarin for starting the draft in #46163. I think the general direction makes sense: keeping the default hard semantic map while optionally exposing the pre-final tensor would already make downstream experimentation much easier.

I have one small question about query-based segmentation models. For dense models, the returned tensor seems straightforward: the resized `outputs.logits` before `argmax`. For MaskFormer, the draft follows the existing post-processing path:

```python
masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
masks_probs = masks_queries_logits.sigmoid()
segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
```

If $L^{cls}$ denotes `class_queries_logits` and $L^{mask}$ denotes `masks_queries_logits`, then for class $c$, query $q$, and pixel $(h, w)$, the `einsum` can be read as:

$$S_{c,h,w} = \sum_{q=1}^{Q} softmax(L^{cls}_{q,:})_c \cdot sigmoid(L^{mask}_{q,h,w})$$

This is useful because it is the tensor used before the final `argmax`. My only uncertainty is how downstream libraries should interpret it: unlike dense `outputs.logits`, this tensor is already built from `softmax` class scores and `sigmoid` mask scores, then combined across queries.

Do you think downstream code should treat this as pre-`argmax` semantic scores, as probabilities, or simply as the model-specific tensor before the hard semantic map is produced? I do not have a strong API suggestion here; I mainly want to avoid assuming the wrong semantics for query-based models when adapting methods such as RankSEG.


Good points! The naming isn't final yet and you are completely right that the returned values are not always logits. I'm considering renaming to `return_scores` instead. I think for consistency with the final segmentation map it makes sense to return the values right before the argmax and not earlier values like `mask_logits` or `class_logits`. There is also at least one model (SegGPT) that uses argmin instead of argmax:
https://github.com/huggingface/transformers/blob/0588858f54c8c79d28497d3ad6eac3417b716c49/src/transformers/models/seggpt/image_processing_seggpt.py#L315

Not yet 100% sure what the best way to handle this is. 

Thanks @guarin, that makes sense to me. `return_scores` sounds safer than `return_logits` for this API.

Thanks again for pushing this forward.

Quick sync from the other side 🤗

Head on #46163 is `e11512d`, `return_segmentation_scores` / `segmentation_scores: torch.Tensor | None` / `SemanticSegmentationPostProcessorOutput` over `PostProcessorOutput(UserDict)`. Ready-for-review, molbap approved on Jun 8, not merged yet.

I opened https://github.com/rankseg/rankseg/issues/26 to note down what the RankSEG side needs once #46163 lands, and posted the score-only ask on the PR itself (a paired `return_segmentation: bool = True` flag, so `return_segmentation_scores=True, return_segmentation=False` skips the argmax and returns the dataclass with `segmentation=None`). Will report back here once #46163 lands and we've retested the new API. 🫡


Thanks again for the discussion around #46163 🤗

I saw @guarin's reply there, and I agree this should not expand the scope of that PR. Exposing `segmentation_scores` already feels like the useful first step.

I wanted to move the RankSEG follow-up back here, since it is more about the integration boundary. Would also be happy to hear thoughts from @Rocketknight1 and @molbap if you have bandwidth 🙏

For RankSEG, the point is not really saving the small cost of an extra `argmax`. The hard segmentation map is not what RankSEG consumes. What it needs is the restored semantic score tensor before the final hard decision.

The split I have in mind is roughly:

```text
Transformers processor:
raw model outputs -> restored semantic scores

Optional RankSEG integration:
restored semantic scores -> hard segmentation map
```

One possible shape would be an isolated optional integration function, rather than adding RankSEG-specific logic to every processor:

```python
from transformers.integrations.rankseg import decode_semantic_segmentation

outputs = processor.post_process_semantic_segmentation(
    model_outputs,
    target_sizes=target_sizes,
    # possible return_segmentation=False
    return_segmentation_scores=True,
)

rankseg_mask = decode_semantic_segmentation(
    outputs,
    metric="dice",
    solver="RMA",
)
```

In that version, `rankseg` would stay optional, e.g. installed through something like `transformers[rankseg]`.

So I wanted to ask what boundary would fit Transformers best:

1. docs-only, with RankSEG documenting the integration externally;
2. optional integration under something like `transformers.integrations.rankseg.decode_semantic_segmentation`;
3. processor-level API.

My intuition is option 2, but I do not want to assume this belongs in Transformers before asking.


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Returning segmentation scores from post-processing was added in https://github.com/huggingface/transformers/pull/46163