Distribution-aware Margin Calibration for Medical Image Segmentation

TL;DR

This paper introduces a distribution-aware margin calibration technique to improve the generalization and performance of mean IoU optimization in medical image segmentation, demonstrating significant empirical gains.

Contribution

It proposes a novel margin calibration method with a theoretical lower bound to enhance mIoU generalization in medical image segmentation.

Findings

Substantial IoU score improvements on two datasets

Better generalization of mIoU over data distribution

Effective in deep segmentation models

Abstract

The Jaccard index, also known as Intersection-over-Union (IoU score), is one of the most critical evaluation metrics in medical image segmentation. However, directly optimizing the mean IoU (mIoU) score over multiple objective classes is an open problem. Although some algorithms have been proposed to optimize its surrogates, there is no guarantee provided for their generalization ability. In this paper, we present a novel data-distribution-aware margin calibration method for a better generalization of the mIoU over the whole data-distribution, underpinned by a rigid lower bound. This scheme ensures a better segmentation performance in terms of IoU scores in practice. We evaluate the effectiveness of the proposed margin calibration method on two medical image segmentation datasets, showing substantial improvements of IoU scores over other learning schemes using deep segmentation models.