RankSEG: A Consistent Ranking-based Framework for Segmentation
Authors: Ben Dai, Chunlin Li; Published in Journal of Machine Learning Research, 24(224):1−50, 2023.
Abstract
Segmentation is an important field in computer vision and natural language processing, where labels are assigned to pixels or features in order to extract regions of interest from images or text. The performance of segmentation is typically evaluated using metrics such as Dice and IoU, which measure the overlap between the predicted segmentation and the ground truth. In this paper, we establish a theoretical foundation for segmentation based on the Dice and IoU metrics. We introduce the concepts of Bayes rule and Dice-/IoU-calibration, which are analogous to classification-calibration or Fisher consistency in classification. We prove that the existing thresholding-based frameworks with operating losses are not consistent with respect to the Dice and IoU metrics, and can lead to suboptimal solutions. To address this issue, we propose a novel consistent ranking-based framework called RankDice/RankIoU, inspired by the plug-in rules of the Bayes segmentation rule. We develop three numerical algorithms with GPU parallel execution to implement the proposed framework for large-scale and high-dimensional segmentation. We also study the statistical properties of the proposed framework, showing that it is Dice-/IoU-calibrated and providing excess risk bounds and the rate of convergence. We demonstrate the numerical effectiveness of RankDice/mRankDice using various simulated examples and datasets such as Fine-annotated CityScapes, Pascal VOC, and Kvasir-SEG with state-of-the-art deep learning architectures. The Python module and source code for RankSEG are available on GitHub at (https://github.com/statmlben/rankseg).
[Abstract]