Stratification for Image Segmentation

Label-aware dataset splitting for image segmentation, giving more representative test sets and more stable model evaluation.

Random splits of segmentation datasets often produce unrepresentative test sets, which biases evaluation. Stratified sampling solves this for classification, but segmentation is harder: every image carries several labels, usually with strong class imbalance.

We introduce Iterative Pixel Stratification (IPS), a simple label-aware sampling method for segmentation, and Wasserstein-Driven Evolutionary Stratification (WDES), a genetic algorithm that minimises the Wasserstein distance between the label distributions of the splits.

Ideal and random splits of PascalVOC plane images
PascalVOC images of the class "plane". An ideal split balances the classes across folds; random splitting does not guarantee this.

Highlights

  • WDES is provably globally optimal given enough generations.
  • Using new statistical heterogeneity metrics, WDES consistently produces more representative splits than random sampling.
  • Across street scenes, medical imaging and satellite imagery, WDES lowers performance variance and improves model evaluation.
  • The gains are largest on small, imbalanced and low-diversity datasets, where conventional splitting is most prone to bias.
Relative improvement of WDES over random splitting
Relative improvement of WDES over random splitting. The gap shrinks with more classes, but WDES keeps an edge on datasets with high class ubiquity.
Stratification method rankings by performance variance
Method rankings by performance variance. WDES ranks highest in low-entropy settings; random splitting catches up as entropy grows.
Runtime of random splitting, IPS and WDES
Runtime of random splitting, IPS and WDES across all datasets.

Outputs

Stratify or Die poster
NeurIPS 2025 poster. Click to open the PDF.