An algorithm that directly optimizes Dice/IoU

[ZixunWang]

Greetings!

We would like to share our recent project on rankseg, a segmentation algorithm that directly optimizes Dice and IoU metrics. It works with any pre-trained segmentation model by only replacing the inference/prediction step. Please feel free to play with it by very minor modification:

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', solver='RMA')

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

If you are interested, please also check out our related papers:

• [JMLR] RankSEG: A Consistent Ranking-based Framework for Segmentation
• [NeurIPS 2025] RankSEG-RMA: An Efficient Segmentation Algorithm via Reciprocal Moment Approximation