Discussion: optional RankSEG-style decoding support for nnU-Net inference

Hi nnU-Net maintainers,

I would like to ask whether an optional RankSEG-style decoder could be worth discussing for the nnU-Net inference/export stage. I may be missing nnU-Net-specific constraints, so this is only a discussion issue, not a concrete API request.

I ran a small smoke test using RankSEG on top of an official TotalSegmentator checkpoint in nnU-Net v2 format. The result is not a benchmark, but it looked encouraging enough to ask whether this direction is worth discussing.

Small Smoke Test

I tested this externally with an official TotalSegmentator organs checkpoint in nnU-Net v2 format:

• checkpoint: Dataset291_TotalSegmentator_part1_organs_1559subj.zip
  • https://github.com/wasserth/TotalSegmentator/releases/download/v2.0.0-weights/Dataset291_TotalSegmentator_part1_organs_1559subj.zip
• dataset: TotalSegmentator organs official-test cases
  • https://huggingface.co/datasets/Angelou0516/totalsegmentator-organs
• subset: 20 smallest official-test cases by remote ct.nii.gz + organs_label.nii.gz file size
• no retraining or checkpoint modification
• baseline: nnU-Net probability maps decoded with voxel-wise argmax
• RankSEG: metric="dice", solver="RMA", output_mode="multiclass"

The metric code is standalone. For this non-region, no-ignore-label demo, it uses the same basic foreground-label averaging convention as nnU-Net-style evaluation: compute per case/label Dice and IoU, average each label across cases with np.nanmean, then average foreground label means.

This is not a benchmark. The gain is small, the subset is size-selected, and some labels declined, likely due to the small sample size. The point is only that the method can be inserted after nnU-Net probability prediction without retraining, and there is enough signal to ask whether the direction is reasonable.

Example visualizations:

RankSEG References

• RankSEG: A Consistent Ranking-based Framework for Segmentation, Dai and Li, JMLR 2023: https://www.jmlr.org/papers/v24/22-0712.html
• RankSEG-RMA: An Efficient Segmentation Algorithm via Reciprocal Moment Approximation, Wang and Dai: https://arxiv.org/abs/2510.15362
• RankSEG package/source: https://pypi.org/project/rankseg/ and https://github.com/rankseg/rankseg

cc @ZixunWang and @statmlben, with thanks for making the RankSEG work and implementation available.

Why this may fit nnU-Net

nnU-Net treats segmentation quality as an end-to-end, dataset-dependent pipeline problem: preprocessing, architecture, training, inference, ensembling, and postprocessing are evaluated choices rather than universally fixed assumptions.

From that perspective, voxel-wise argmax is a strong and efficient default, but it is still one probability-to-label decision rule. RankSEG-style decoding could be treated as an optional inference-time candidate.

Possible Integration Point

If nnU-Net maintainers think this direction is worth exploring, one possible integration point could be the stage where nnU-Net converts model outputs into the final hard segmentation. It would not affect nnU-Net planning, preprocessing, training, checkpoints, or architectures.

From reading the current code, the relevant stage may be around convert_predicted_logits_to_segmentation_with_correct_shape in nnunetv2/inference/export_prediction.py, where model outputs are converted into the final segmentation. But I am not assuming this is the right API boundary.

nnunetv2/inference/export_prediction.py
convert_predicted_logits_to_segmentation_with_correct_shape(...)

Closing Note

I would mainly appreciate feedback on whether an optional metric-aware decoding mode is compatible with nnU-Net's design goals.

If the maintainers think there is a path forward, we 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.