Random splitting of datasets in image segmentation often leads to unrepresentative test sets, resulting in biased evaluations and poor model generalization. While stratified sampling has proven effective for addressing label distribution imbalance in classification tasks, extending these ideas to segmentation remains challenging due to the multi-label structure and class imbalance typically present in such data. Building on existing stratification concepts, we introduce Iterative Pixel Stratification (IPS), a straightforward, label-aware sampling method tailored for segmentation tasks. Additionally, we present Wasserstein-Driven Evolutionary Stratification (WDES), a novel genetic algorithm designed to minimize the Wasserstein distance, thereby optimizing the similarity of label distributions across dataset splits. We prove that WDES is globally optimal given enough generations. Using newly proposed statistical heterogeneity metrics, we evaluate both methods against random sampling and find that WDES consistently produces more representative splits. Applying WDES across diverse segmentation tasks, including street scenes, medical imaging, and satellite imagery, leads to lower performance variance and improved model evaluation. Our results also highlight the particular value of WDES in handling small, imbalanced, and low-diversity datasets, where conventional splitting strategies are most prone to bias.

翻译：图像分割中数据集的随机划分通常导致测试集缺乏代表性，从而引发评估偏差和模型泛化能力不足。虽然分层抽样已被证明能有效解决分类任务中的标签分布不平衡问题，但由于分割数据通常具有多标签结构和类别不平衡特性，将这些思想扩展到分割任务仍面临挑战。基于现有分层概念，我们提出了迭代像素分层（IPS）——一种专为分割任务设计的简单、标签感知的抽样方法。此外，我们提出了瓦瑟斯坦驱动进化分层（WDES）——一种新颖的遗传算法，旨在最小化瓦瑟斯坦距离，从而优化数据集划分间标签分布的相似性。我们证明了WDES在足够多代演化后具有全局最优性。通过新提出的统计异质性度量指标，我们将两种方法与随机抽样进行对比评估，发现WDES始终能产生更具代表性的划分。将WDES应用于街景、医学影像和卫星图像等多种分割任务中，有效降低了性能方差并改善了模型评估效果。我们的结果还突显了WDES在处理小型、不平衡及低多样性数据集时的特殊价值，这些场景下传统划分策略最容易产生偏差。