rankseg 0.0.5

RankSEG: A Statistically Consistent Segmentation Prediction Solver for Dice and IoU Metrics Optimization

🧩 RankSEG

Boost Segmentation Performance Instantly via Direct Dice/IoU Post-Optimization

Quick Start | Official Integrations | Key Features | Benchmarks | Citation

---

RankSEG is a plug-and-play post-processing module that boosts segmentation performance (Dice/IoU) during inference. It works with ANY pre-trained probabilistic segmentation model (SAM, DeepLab, SegFormer, etc.) without any retraining or fine-tuning.

Explore RankSEG by reading our documentation.

If RankSEG improves your segmentation workflow, please consider starring the repo: https://github.com/rankseg/rankseg

🌟 Why RankSEG?

Conventional methods use argmax or fixed thresholding, which are not theoretically optimized for non-decomposable metrics like Dice or IoU. RankSEG bridges this gap by directly optimizing the target metric, yielding "free" performance gains.

Demo: RankSEG vs. Argmax on Segformer ADE20k

⚡ Quick Start

RankSEG is designed to drop into an existing PyTorch segmentation pipeline with just a few lines of code.

1. Installation

pip install -U rankseg

2. Basic Usage

PyTorch Native Flow

from rankseg import RankSEG
import torch.nn.functional as F

# 1. Initialize RankSEG with the official default configuration
rankseg = RankSEG(metric="dice", solver="RMA", output_mode="multiclass")

# 2. Get probability output from YOUR model
# probs: (batch_size, num_classes, *image_shape)
probs = F.softmax(model_logits, dim=1)

# 3. Get optimized predictions
preds = rankseg(probs)

You can also use the functional API for one-off prediction:

from rankseg.functional import rankseg

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

💡 Try it now:

Official PyTorch integration: Docs · Example

Hugging Face semantic segmentation integration: Notebook · Colab

SAM family integration: Notebook · Colab

🔌 Official Integrations

These are the maintained integration entry points documented by this repository.

Path	Status	Entry
PyTorch Native	Ready	Docs · Example
Hugging Face semantic segmentation	Ready	from rankseg.integration import transformers -> transformers.postprocess / transformers.restore_semantic_probs · Docs · Example
SAM family	Ready	from rankseg.integration import sam -> sam.Sam1 / sam.Sam2 / sam.Sam3 · Docs · Notebook

🌐 External Integrations

These integrations already exist, but are currently maintained outside the main repository.

Integration	Status	Entry
PaddleSeg	External	Docs · Branch · Docker

✨ Key Features

• 🚀 Performance Boost: Consistently improves mIoU/mDice scores over standard argmax.
• 🔌 Zero Effort: Compatible with any PyTorch model. No retraining, no fine-tuning.
• 🆓 Training-Free: Purely post-processing. Works with frozen weights.
• ⚡ Real-time Inference: Efficient RMA (Reciprocal Moment Approximation) solver.
• 🧩 Versatile: Supports semantic (multi-class) and binary (multi-label) tasks.

📊 Benchmarks

RankSEG delivers consistent gains across various architectures and datasets without touching a single weight.

Model	Dataset	mIoU (Argmax)	mIoU (RankSEG)	Gain
DeepLabV3+	PASCAL VOC	77.25%	78.14%	+0.89%
SegFormer	PASCAL VOC	77.57%	78.59%	+1.02%
UPerNet	PASCAL VOC	79.52%	80.31%	+0.79%
SegFormer	ADE20K	40.00%	40.82%	+0.82%
UPerNet	ADE20K	42.86%	43.84%	+0.98%

Detailed results available in our NeurIPS 2025 paper.

🧪 Additional Demos

Framework	Task	Quick Start
SAM family	SAM1, SAM2, SAM3 masks
Hugging Face	Interactive Demo

🔗 Citation

If you use RankSEG in your research, please cite our papers:

• Dai, B., & Li, C. (2023). RankSEG: A Consistent Ranking-based Framework for Segmentation. Journal of Machine Learning Research, 24(224), 1-50. [link]
• Wang, Z., & Dai, B. (2025). RankSEG-RMA: An Efficient Segmentation Algorithm via Reciprocal Moment Approximation. Advances in Neural Information Processing Systems (NeurIPS 2025). [link]

@article{dai2023rankseg,
  title={RankSEG: A Consistent Ranking-based Framework for Segmentation},
  author={Dai, Ben and Li, Chunlin},
  journal={Journal of Machine Learning Research},
  volume={24},
  number={224},
  pages={1--50},
  url={https://www.jmlr.org/papers/v24/22-0712.html},
  year={2023}
}

@inproceedings{wang2025rankseg,
  title={RankSEG-RMA: An Efficient Segmentation Algorithm via Reciprocal Moment Approximation},
  author={Wang, Zixun and Dai, Ben},
  booktitle={Advances in Neural Information Processing Systems},
  url={https://arxiv.org/abs/2510.15362},
  year={2025}
}

---

Star us on GitHub if RankSEG helps your project! ⭐