Make `argmax` in `post_process_semantic_segmentation` optional
Author: simonreise
Pull request: no
Labels: name: Feature request

### Feature request

Now `post_process_semantic_segmentation` function generates [N, H, W] resized segmentation maps, but sometimes user might need [N, C, H, W] resized class probabilities.

Maybe it's worth adding a bool optional argument (named e.g. `class_proba`) to `post_process_semantic_segmentation` that will control if `argmax` should or shouldn't be performed.

### Motivation

There might be cases when user need not resized segmentation maps of [N, H, W] shape, but resized class probabilities of [N, C, H, W] shape.

Calculating ROC AUC metric or a custom segmentation loss (e. g. one of [SMP Losses](https://smp.readthedocs.io/en/latest/losses.html)) might be the case.

#### More detailed reasoning

As far as I understand (mb I am wrong) applying `outputs.logits.softmax(dim=1)` to raw model outputs won't work for every model, because some models (like DETR or MaskFormer) produce `class_queries_logits` and `mask_queries_logits` that have to be converted to something more segmentation-map-alike first. Also, many models need their outputs to be resized that is also done in `post_process_semantic_segmentation`.

In my case I wanted my geospatial semantic segmentation library to be able to get class probabilities from any HF Transformers semantic segmentation model output and I ended copying post-processing code for DETR and ConditionalDETR, for MaskFormer and Mask2Former, for OneFormer and for other segmentation models separately (because they process raw model outputs in different ways) and merging them into one function (that looked extremely ugly).

<details>
<summary>Warning bad code</summary>

```py
def process_output(self, pred, x):
        if (
            isinstance(pred, transformers.modeling_outputs.SemanticSegmenterOutput) 
            or isinstance(pred, transformers.models.clipseg.modeling_clipseg.CLIPSegImageSegmentationOutput)
        ):
            pred = pred.logits
            if pred.shape[2:4] != x.shape[2:4]:
                pred = torch.nn.functional.interpolate(
                    pred,
                    size=x.shape[2:4],
                    mode="bilinear",
                    align_corners=False,
                )
            pred = pred.softmax(dim=1)
        elif (
            isinstance(pred, transformers.models.conditional_detr.modeling_conditional_detr.ConditionalDetrSegmentationOutput) 
            or isinstance(pred, transformers.models.detr.modeling_detr.DetrSegmentationOutput)
        ):
            class_queries_logits = pred.logits  # [batch_size, num_queries, num_classes+1]
            masks_queries_logits = pred.pred_masks  # [batch_size, num_queries, height, width]
            
            # Remove the null class `[..., :-1]`
            # Detr treats 0 as separate class, conditionaldetr treats it as nullclass
            if isinstance(pred, transformers.models.detr.modeling_detr.DetrSegmentationOutput):
                masks_classes = class_queries_logits.softmax(dim = -1)[..., :-1]
            else:
                masks_classes = class_queries_logits.softmax(dim=-1)
            masks_probs = masks_queries_logits.sigmoid()  # [batch_size, num_queries, height, width]

            # Semantic segmentation logits of shape (batch_size, num_classes, height, width)
            pred = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)

            if pred.shape[2:4] != x.shape[2:4]:
                pred = torch.nn.functional.interpolate(
                    pred,
                    size=x.shape[2:4],
                    mode="bilinear",
                    align_corners=False,
                )
            pred = pred.softmax(dim=1)
        elif (
            isinstance(pred, transformers.models.mask2former.modeling_mask2former.Mask2FormerForUniversalSegmentationOutput) 
            or isinstance(pred, transformers.models.maskformer.modeling_maskformer.MaskFormerForInstanceSegmentationOutput)
        ):
            class_queries_logits = pred.class_queries_logits  # [batch_size, num_queries, num_classes+1]
            masks_queries_logits = pred.masks_queries_logits  # [batch_size, num_queries, height, width]
            masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
            masks_probs = masks_queries_logits.sigmoid()  # [batch_size, num_queries, height, width]
            
            # Semantic segmentation logits of shape (batch_size, num_classes, height, width)
            pred = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
            if pred.shape[2:4] != x.shape[2:4]:
                pred = torch.nn.functional.interpolate(
                    pred,
                    size=x.shape[2:4],
                    mode="bilinear",
                    align_corners=False,
                )
            pred = pred.softmax(dim=1)
        elif isinstance(pred, transformers.models.oneformer.modeling_oneformer.OneFormerForUniversalSegmentationOutput):
            class_queries_logits = pred.class_queries_logits  # [batch_size, num_queries, num_classes+1]
            masks_queries_logits = pred.masks_queries_logits  # [batch_size, num_queries, height, width]
            masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
            masks_probs = masks_queries_logits.sigmoid()  # [batch_size, num_queries, height, width]

            # Semantic segmentation logits of shape (batch_size, num_classes, height, width)
            pred = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
            if pred.shape[2:4] != x.shape[2:4]:
                pred = torch.nn.functional.interpolate(
                    pred,
                    size=x.shape[2:4],
                    mode="bilinear",
                    align_corners=False,
                )        
            pred = pred.softmax(dim=1)
        elif isinstance(pred, collections.OrderedDict):
            pred = pred['out']
        return pred
```
</details>

I think user should be able to get class probabilities in a more intuitive way. That's why I opened an original issue with a feature request. 

If the requested feature is implemented, the user would be able to get class probabilities by just calling `probabilities = processor.post_processs_semantic_segmentation(predictions, target_sizes=size, class_proba=True)` with every model.

### Your contribution

Implementing it doesn't look hard, but it would require to change every ImageProcessor and FastImageProcessor that support semantic segmentation

`class_proba` is optional and False by default for backward compatibility.

```py
def post_process_semantic_segmentation(
  self, 
  outputs, 
  target_sizes: Optional[List[Tuple]] = None,
  class_proba: Optional[bool] = False,
):
"""
...
class_proba (`bool`, *optional*, defaults to `False`):
  Whether to keep class probabilities.
  Will return (N, H, W) segmentation maps if False and (N, C, H, W) class probabilities maps if True.
"""
...
...
...
# Old code:
# semantic_map = resized_logits[0].argmax(dim=0)
# New code:
if class_proba:
  semantic_map = resized_logits[0].softmax(dim=0)
else:
  semantic_map = resized_logits[0].argmax(dim=0)
...
...
...
# Old code:
# semantic_segmentation = logits.argmax(dim=1)
# New code:
if class_proba:
  semantic_segmentation = logits.softmax(dim=1)
else:
  semantic_segmentation = logits.argmax(dim=1)
```

UPD: looks like most of the custom metrics or losses (like Torchmetrics AUROC or SMP Jaccard Loss) support logits as inputs, so we do not need to convert logits to probabilities. But we anyway need to resize raw model outputs or convert class and mask queries, so the proposed feature anyway will be useful.

Also, we can give more control to user by implementing something like:

```py
def post_process_semantic_segmentation(
  self, 
  outputs, 
  target_sizes: Optional[List[Tuple]] = None,
  output_format: Optional[Literal["segmentation_maps", "logits", "class_proba"]] = "segmentation_maps",
):
"""
...
output_format(`str`, *optional*, defaults to `"segmentation_maps"`):
    Output data format.
    "segmentation_maps": will return (N, H, W) segmentation maps. 
    "class_proba": will return (N, C, H, W) class probabilities maps. Class probabilities are in range [0, 1] and sum to 1.
    "logits": will return (N, C, H, W) raw logits.
"""
...
if output_format == "segmentation_maps":
    semantic_map = resized_logits[0].argmax(dim=0)
elif output_format == "class_proba":
    semantic_map = resized_logits[0].softmax(dim=0)
elif output_format == "logits":
    semantic_map = resized_logits[0]
else:
    raise ValueError("Unknown output format")

Comments:
cc @qubvel @nielsrogge

Pls check #37904

I wanted to add one data point to this discussion, not so much as a concrete proposal for how `transformers` should solve it internally, but rather as support for the underlying user need behind this issue that @simonreise raised and that @qubvel and @nielsrogge have already discussed from the library/API side.

I agree that exposing something more flexible than the final `argmax` semantic map can be meaningful in practice.

In my case, this came up while I was experimenting with alternative semantic post-processing methods, in particular [RankSEG](https://github.com/rankseg/rankseg). I also want to take the opportunity to thank @statmlben and @ZixunWang here, since this line of work is exactly what made me feel that the demand in this issue is real and not just a convenience request.

RankSEG is a plug-and-play post-processing method for probabilistic segmentation outputs:

- repo: <https://github.com/rankseg/rankseg>
- docs: <https://rankseg.readthedocs.io/en/latest/>
- Hugging Face Space: <https://huggingface.co/spaces/statmlben/rankseg>
- JMLR paper: <https://www.jmlr.org/papers/v24/22-0712.html>
- NeurIPS 2025 paper: <https://openreview.net/pdf?id=4tRMm1JJhw>

What RankSEG made very visible to me is the practical need behind this issue: for some downstream methods, the final hard semantic map is already too late. What is useful is access to semantic scores / probabilities before the last post-processing decision is made.

Part of why I found this compelling is that RankSEG is not just a cosmetic alternative to `argmax`. The method is designed to optimize segmentation metrics such as Dice/IoU more directly at inference time, and in the public RankSEG materials there are consistent gains reported over standard `argmax` post-processing across several models and datasets, without retraining or changing model weights. For example, the README reports improvements on setups such as SegFormer and UPerNet on PASCAL VOC and ADE20K. That makes the quality of the pre-final semantic output practically important, not just theoretically interesting.

I have been exploring this direction in a prototype branch here:

- Branch: <https://github.com/Leev1s/transformers/tree/leev1s/semantic-rankseg>
- Current implementation draft: <https://github.com/Leev1s/transformers/blob/leev1s/semantic-rankseg/src/transformers/pipelines/image_segmentation.py>
- Current HEAD commit: <https://github.com/Leev1s/transformers/commit/49f24638349c6072da02c60f2501c0000372c863>

Rather, my point is that this issue reflects a real need:

- for some workflows, users do not only want the final semantic map
- they want an intermediate representation that is still usable by downstream post-processing
- this is especially relevant when semantic predictions are not just `outputs.logits.softmax(dim=1)`, but have to be assembled first from query logits and mask logits

For example, in the prototype I first prepare semantic probabilities and only then apply an alternative semantic post-processing step:

```python
if getattr(model_outputs, "logits", None) is not None and getattr(model_outputs, "pred_masks", None) is None:
    probs = model_outputs.logits.float().softmax(dim=1)
    return [probs[i] for i in range(probs.shape[0])]
```

For query-based segmentation models, I reconstruct semantic scores before the final post-processing:

```python
class_queries_logits = getattr(model_outputs, "class_queries_logits", None)
masks_queries_logits = getattr(model_outputs, "masks_queries_logits", None)

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

return [
    scores / scores.sum(dim=0, keepdim=True).clamp_min(1e-12) for scores in segmentation_scores
]
```

Then that representation can be consumed by another semantic post-processing method:

```python
probs = self._get_semantic_segmentation_probs(model_outputs)[0]
segmentation = rankseg.predict(probs.unsqueeze(0))[0].to(torch.long)
```

From my perspective, the interesting part of this issue is not necessarily whether `post_process_semantic_segmentation(...)` itself should gain a `class_proba` flag.

The more general point is that the current semantic post-processing path is still largely managed inside model-specific code, and that makes advanced downstream post-processing harder than it needs to be. This is understandable, because `transformers` is a model framework, not a library whose goal is to ship every possible task-specific post-processing strategy.

But I do think it would be valuable if the library had a clearer convention, or at least a more flexible way, for models / processors to expose semantic outputs before the last hard post-processing step. That would make it easier for downstream users and external libraries to experiment with their own post-processing methods without having to re-implement model-specific conversion logic.

That is the main reason I think the demand behind this issue is valid.

So my take is:

- not necessarily "please add this exact API"
- but rather "please recognize that exposing pre-final semantic outputs has real downstream value"

And from that angle, this issue is relevant beyond just one specific post-processing feature.

For reference, my RankSEG-related exploration is here again:

- <https://github.com/Leev1s/transformers/tree/leev1s/semantic-rankseg>
- <https://github.com/rankseg/rankseg>

I would be very interested in hearing whether others also see value in making this stage easier to access, even if the library does not want to standardize one particular post-processing algorithm.

@yonigozlan @molbap

#46163 and #46139 discuss and implement the behavior I was asking for. IMHO this issue can be closed, and, probably #37904 can be closed too