Discussion: optional RankSEG-style decoding support in MONAI post-processing #8857

Hi MONAI maintainers,

I would like to ask for your feedback on whether RankSEG-style decoding could be a good fit for MONAI post-processing. I ran a small smoke test using RankSEG on top of a pretrained MONAI bundle, and the result looked encouraging, especially for the tumor class. This issue is meant as an initial discussion rather than a request to merge a specific API.

Preliminary Smoke Test

I used the pretrained MONAI pancreas_ct_dints_segmentation bundle on 6 MSD Pancreas cases. This is not intended as a full benchmark, and the sample size is small. The goal was only to check whether RankSEG can be inserted into a standard MONAI bundle workflow and whether the result is worth further discussion.

Setup:

• MONAI: 1.5.2
• PyTorch: 2.12.0+cu130
• Bundle: pancreas_ct_dints_segmentation
• Task: MSD Pancreas, 3 classes: background, pancreas, tumor
• No retraining
• No checkpoint modification
• Post-processing only

Results:

In this small test, the largest visible change was on the tumor class, where Dice improved from 0.6188 with argmax to 0.6654 with RankSEG.

Example visualization:

The visualization focuses on tumor segmentation. Error maps use white for true positive, red for false positive, and cyan for false negative.

Broader Evidence From RankSEG Papers

The MONAI experiment above is only a small smoke test for this possible integration point. RankSEG has also been studied more systematically in two papers:

• RankSEG: A Consistent Ranking-based Framework for Segmentation,
  Dai and Li, JMLR 2023.
• RankSEG-RMA: An Efficient Segmentation Algorithm via Reciprocal Moment Approximation,
  Wang and Dai, NeurIPS 2025.

These papers are the main evidence base for RankSEG. The MONAI pancreas bundle result above is only meant to show that the same idea can be inserted into a standard MONAI bundle workflow without retraining or checkpoint changes.

Possible Integration Direction

RankSEG is a training-free segmentation decoding method that converts model probability maps into discrete segmentation masks with respect to a target metric such as Dice.

If MONAI maintainers think this direction is worth exploring, one possible integration point could be the same stage where many MONAI pipelines currently use AsDiscrete(argmax=True). The intended scope would be limited to post-processing; it would not affect MONAI model training, bundle loading, preprocessing, or inference APIs.

Current common pattern:

from monai.transforms import AsDiscrete

mask = AsDiscrete(argmax=True)(probs).squeeze(0).long()

Temporary workaround outside MONAI:

from rankseg import RankSEG

rankseg = RankSEG(
    metric="dice",
    solver="RMA",
    output_mode="multiclass",
)

mask = rankseg.predict(probs.unsqueeze(0)).squeeze(0).long()

This works today, but it requires users to step outside the standard MONAI post-processing API. If maintainers think it fits the library, one possible long-term direction would be to make this available through AsDiscrete, for example as an optional decoding mode, so users can keep the same MONAI transform entry point while choosing between plain argmax and RankSEG-style metric-aware decoding.

Here probs is the softmax probability tensor from a MONAI segmentation model. For a 3D multiclass case, this is typically shaped like (C, H, W, D).

Why This May Fit MONAI

RankSEG does not require retraining and does not require a new checkpoint. It uses the probability maps from an already trained segmentation model and produces the final mask as a post-processing step.

This may make it a possible fit for MONAI's existing post-processing transform layer:

model logits
-> Activations(softmax=True)
-> AsDiscrete(..., rankseg=True) or RankSEG / RankSEGd
-> metric computation or SaveImage

This might be useful for bundle inference workflows, where users may want to try metric-aware decoding without changing the model or retraining.

Closing Note

We understand that maintaining a large and widely used project like MONAI requires careful review of API design, dependencies, tests, documentation, and long-term maintenance burden. I do not want to assume that this is the right fit for MONAI, but I would appreciate any feedback on whether this direction is worth exploring.

If the maintainers think there is a path forward, I and our team would be happy to help with additional tests, API revisions, documentation, examples, or any engineering work needed to make the contribution easier to review and maintain.