Feat: Add class_proba option to semantic segmentation post-processing
Author: demoncoder-crypto
Pull request: yes

Addresses #37715, this is my first implementation, if my direction is correct please do let me know I will iteratively fix everything 

Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

cc @qubvel @nielsrogge

Hey @demoncoder-crypto! Thanks for opening a PR 🤗 In case you want to obtain class probabilities for image pixels, you can always apply sigmoid/softmax activation function to the raw model outputs (avoiding applying post_process_semantic_segmentation). For example:

```python
with torch.no_grad():
    outputs = model(**inputs)

probabilities = outputs.logits.softmax(dim=1)
```

The purpose of `post_process_semantic_segmentation` is to provide you with a segmentation map, so I'm not sure we should add another option for this method. Thank you!

@qubvel Thanks for your reply! But 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. 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.

P. S. Also thanks for noticing that instead of removing argmax we should replace it with softmax in case we need class probabilities.

Hi everyone, thanks for the work on this PR — I'm also interested in obtaining pixel-level class probabilities during post-processing for semantic segmentation.

I noticed that the current implementation does not compute true pixel-level probability distributions. In particular, this line:
```segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)```
produces a tensor of shape [b, c, h, w] whose values are not guaranteed to form a valid probability distribution over classes at each pixel. Specifically, the following is False in general:
 ```torch.isclose(segmentation.sum(1), torch.tensor(1.)).all()```.
This happens because each pixel-class score is a weighted sum over all queries, combining the objectness scores at a pixel with the class probabilities per query — but there's no guarantee this will sum to 1 across classes.

In my use case, I need the output to be a valid per-pixel class probability distribution (i.e., summing to 1 over classes), so I'm currently normalizing the result like this (and getting quite acceptable results in pratice):
```
segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
segmentation = segmentation / segmentation.sum(1, keepdim=True)
```
An alternative might be to apply softmax() over the class dimension after the einsum, but this tends to overly smooth the predictions, reducing confidence in the most likely classes.

Just wanted to share this in case it's helpful, and to see if others have encountered similar needs or solutions.