RankSEG: Boost Semantic Segmentation Accuracy with Just Three Lines of Code

In semantic segmentation the common inference step applies a fixed threshold or argmax on the pixel‑wise probability map. This practice optimises per‑pixel accuracy but does not guarantee optimal region‑level overlap metrics such as Dice or IoU, leading to sub‑optimal segmentation quality.

Limitations of Threshold/Argmax

Because the loss is defined per pixel, the optimal mask for Dice may require labeling a pixel with probability < 0.5 as foreground. A simple two‑pixel example shows that the fixed‑threshold or argmax prediction yields a lower Dice score than the mask that maximises the expected Dice.

Core Theory – RankSEG

RankSEG proves that the Dice‑optimal mask can be obtained by sorting all pixel probabilities in descending order and selecting the top‑k pixels, where k is chosen to maximise the expected Dice coefficient. This reduces the exponential search over all binary masks to a linear scan over possible volumes (k = 0…N), exploiting the ranking property: swapping a higher‑probability pixel into the foreground always increases the expected Dice.

Efficient Approximation – RankSEG‑RMA

Exact computation of the Dice expectation for each candidate k is expensive because it involves a reciprocal‑expectation term. RankSEG‑RMA replaces this term with a Reciprocal Moment Approximation that uses prefix sums of the sorted probabilities. The approximation can be pre‑computed once per image, yielding an overall O(N) complexity where N is the number of pixels.

Multi‑Class Extension

For multi‑label segmentation RankSEG is applied independently per class. For single‑label multi‑class segmentation a greedy strategy is used:

Generate binary masks for each class with RankSEG‑RMA.

Remove overlapping regions, leaving some pixels unassigned.

For each unassigned pixel compute the Dice gain of assigning it to each class.

Assign the pixel to the class with the highest gain.

This retains the benefits of RankSEG while respecting the non‑overlap constraint.

Implementation

from rankseg import RankSEG

# Initialise RankSEG to optimise Dice
rankseg = RankSEG(metric='dice')

# Model probability outputs (batch, num_classes, H, W)
probs = model(images).softmax(dim=1)

# Replace the usual argmax post‑processing
preds = rankseg.predict(probs)

Experimental Results

Extensive experiments on PASCAL VOC, Cityscapes, LiTS and KiTS with various backbone networks show that RankSEG consistently improves Dice/IoU over the baseline argmax. RankSEG‑RMA achieves almost the same accuracy as the exact RankSEG‑BA while being tens of times faster, making it practical for real‑world deployment.

Open‑source implementation and documentation are provided at https://github.com/rankseg/rankseg.