🧩 RankSEG¶
RankSEG is a statistically consistent framework for semantic segmentation that provides plug-and-play modules to improve segmentation results during inference.
RankSEG-based methods are theoretically-grounded segmentation approaches that are statistically consistent with respect to popular segmentation metrics like Dice and IoU. They provide almost guaranteed improved performance over traditional thresholding or argmax segmentation methods.
Note
RankSEG optimizes metrics using a samplewise aggregation: the score is computed per sample and then averaged across the dataset (akin to aggregation_level='samplewise' in TorchMetrics DiceScore). See Metrics for details.
Key Properties¶
🎯 Metric-Optimized |
Directly optimizes for Dice or IoU metrics instead of using generic ad-hoc argmax during inference. |
🔌 Plug-and-Play |
Works with ANY pre-trained prob-outcome segmentation model without retraining |
⚡ Efficient Solvers |
Multiple solver options (BA, TRNA, RMA) for different speed-accuracy trade-offs |
🧩 Flexible Tasks |
Supports both multi-class and multi-label segmentation tasks, whether objects overlap or not. |
✨ Quick Start¶
Install rankseg using pip:
pip install rankseg
Basic usage example:
import torch
import torch.nn.functional as F
from rankseg import RankSEG
# Your pre-trained model's probability output
probs = F.softmax(torch.randn(4, 21, 256, 256), dim=1) # (batch, classes, height, width)
# Create RankSEG predictor optimized for Dice metric
rankseg = RankSEG(metric='dice')
# Get optimized predictions
preds = rankseg(probs)
Functional API for one-off prediction:
from rankseg.functional import rankseg
preds = rankseg(probs, metric='dice')
Why RankSEG?¶
Traditional segmentation methods use argmax or thresholding to convert model outputs to predictions. However, these methods are not optimized for the actual evaluation metrics (Dice, IoU).
Performance Improvements Across Models and Datasets:
RankSEG consistently outperforms standard argmax prediction without any model retraining:
Model |
Dataset |
mIoU (Argmax) |
mIoU (RankSEG) |
mDice (Argmax) |
mDice (RankSEG) |
|---|---|---|---|---|---|
DeepLabV3+ (ResNet101) |
PASCAL VOC |
77.25% |
78.14% ↑0.89% |
82.08% |
83.14% ↑1.06% |
SegFormer (MiT-B4) |
PASCAL VOC |
77.57% |
78.59% ↑1.02% |
82.15% |
83.22% ↑1.07% |
UPerNet (ConvNeXt) |
PASCAL VOC |
79.52% |
80.31% ↑0.79% |
84.11% |
84.98% ↑0.87% |
PSPNet (ResNet101) |
Cityscapes |
65.89% |
66.53% ↑0.64% |
73.55% |
74.28% ↑0.73% |
DeepLabV3+ (ResNet101) |
Cityscapes |
66.17% |
66.68% ↑0.51% |
73.71% |
74.33% ↑0.62% |
UPerNet (ConvNeXt) |
Cityscapes |
68.83% |
69.57% ↑0.74% |
76.08% |
76.97% ↑0.89% |
SegFormer (MiT-B4) |
ADE20K |
40.00% |
40.82% ↑0.82% |
46.50% |
47.57% ↑1.07% |
UPerNet (ConvNeXt) |
ADE20K |
42.86% |
43.84% ↑0.98% |
49.61% |
50.85% ↑1.24% |
CPT (Swin-Large) |
ADE20K |
44.59% |
45.56% ↑0.97% |
51.27% |
52.58% ↑1.31% |
Note
Results from our NeurIPS 2025 paper. RankSEG uses Dice metric with RMA solver.
Learn More¶
Learn how to use RankSEG with your segmentation models
Maintained integration guides for PyTorch, Transformers, and SAM
Notebook tutorials for quickstart, Transformers, SAM, and PaddleSeg
Detailed documentation of all classes and functions
How to cite RankSEG in your research
Source code, issues, and contributions
References¶
Ben Dai and Chunlin Li. Rankseg: a consistent ranking-based framework for segmentation. Journal of Machine Learning Research, 24(224):1–50, 2023.
Zixun Wang and Ben Dai. Rankseg-rma: an efficient segmentation algorithm via reciprocal moment approximation. In Advances in Neural Information Processing Systems. 2025.
A Yu Volkova. A refinement of the central limit theorem for sums of independent random indicators. Theory of Probability & Its Applications, 40(4):791–794, 1996.