Alleviating Class Imbalance in Semi-supervised Multi-organ Segmentation via Balanced Subclass Regularization

TL;DR

This paper introduces BSR-Net, a semi-supervised multi-organ segmentation method that uses balanced subclass regularization to address class imbalance caused by organ size variations, improving segmentation accuracy.

Contribution

The paper proposes a novel two-phase semi-supervised network with balanced subclass generation and auxiliary subclass segmentation to mitigate class imbalance in multi-organ segmentation.

Findings

Outperforms existing methods on MICCAI FLARE 2022 dataset

Effective in handling class imbalance due to organ size variation

Improves segmentation accuracy in semi-supervised learning settings

Abstract

Semi-supervised learning (SSL) has shown notable potential in relieving the heavy demand of dense prediction tasks on large-scale well-annotated datasets, especially for the challenging multi-organ segmentation (MoS). However, the prevailing class-imbalance problem in MoS, caused by the substantial variations in organ size, exacerbates the learning difficulty of the SSL network. To alleviate this issue, we present a two-phase semi-supervised network (BSR-Net) with balanced subclass regularization for MoS. Concretely, in Phase I, we introduce a class-balanced subclass generation strategy based on balanced clustering to effectively generate multiple balanced subclasses from original biased ones according to their pixel proportions. Then, in Phase II, we design an auxiliary subclass segmentation (SCS) task within the multi-task framework of the main MoS task. The SCS task contributes a balanced subclass regularization to the main MoS task and transfers unbiased knowledge to the MoS network, thus alleviating the influence of the class-imbalance problem. Extensive experiments conducted on two publicly available datasets, i.e., the MICCAI FLARE 2022 dataset and the WORD dataset, verify the superior performance of our method compared with other methods.