Marginal loss and exclusion loss for partially supervised multi-organ segmentation

TL;DR

This paper introduces marginal and exclusion loss functions to enable training a multi-organ segmentation network from partially labeled datasets, improving performance without extra computation.

Contribution

The paper proposes novel loss functions tailored for learning from partially labeled multi-organ datasets, addressing label merging and organ disjointness.

Findings

Significant performance improvement on multi-organ segmentation benchmarks.

Effective learning from partially labeled datasets without additional computation.

Applicable to existing segmentation loss functions like cross entropy and Dice loss.

Abstract

Annotating multiple organs in medical images is both costly and time-consuming; therefore, existing multi-organ datasets with labels are often low in sample size and mostly partially labeled, that is, a dataset has a few organs labeled but not all organs. In this paper, we investigate how to learn a single multi-organ segmentation network from a union of such datasets. To this end, we propose two types of novel loss function, particularly designed for this scenario: (i) marginal loss and (ii) exclusion loss. Because the background label for a partially labeled image is, in fact, a `merged' label of all unlabelled organs and `true' background (in the sense of full labels), the probability of this `merged' background label is a marginal probability, summing the relevant probabilities before merging. This marginal probability can be plugged into any existing loss function (such as cross entropy loss, Dice loss, etc.) to form a marginal loss. Leveraging the fact that the organs are non-overlapping, we propose the exclusion loss to gauge the dissimilarity between labeled organs and the estimated segmentation of unlabelled organs. Experiments on a union of five benchmark datasets in multi-organ segmentation of liver, spleen, left and right kidneys, and pancreas demonstrate that using our newly proposed loss functions brings a conspicuous performance improvement for state-of-the-art methods without introducing any extra computation.