MMSeg: A novel multi-task learning framework for class imbalance and label scarcity in medical image segmentation

Feiyang Yang,Xiongfei Li,Bo Wang,Tianyang Zhang,Xinyi Yu,Xingcheng Yi,Rui Zhu
DOI: https://doi.org/10.1016/j.knosys.2024.112835
IF: 8.139
2024-12-11
Knowledge-Based Systems
Abstract:Semi-supervised medical image segmentation can effectively alleviate the high costs associated with obtaining high-quality labels and issues related to data privacy. However, the issue of imbalance in medical images limits the practical application of semi-supervised frameworks. One of the core principles of semi-supervised learning is to infer and predict by learning data distributions. Since unlabeled samples lack explicit class information, both inter-class and intra-class imbalances in the regions of interest can undermine the training core of this framework. To alleviate these issues, we propose a segmentation framework based on multi-task learning, which integrates both self-supervised and semi-supervised learning components within the multi-task framework. In the self-supervised component, we utilize four sub-tasks to extract the internal structure and features of images from unlabeled data, aimed at mitigating the imbalance issues in regions of interest. In the semi-supervised learning component, we employ a segmentation network with a teacher-student model structure that effectively generates high-confidence pseudo-labels from unlabeled data, thereby enhancing the model's ability to recognize minority classes. To maintain the robustness of the model, we introduce a progressive learning strategy and a dynamic weight adjustment mechanism to prevent feature representation bias caused by varying task difficulties. Experimental results indicate that our proposed method achieved the best performance in extensive experiments on two abdominal multi-organ segmentation tasks, outperforming other competitive methods. Codes are available at https://github.com/ReinforceLove/MMSeg .
computer science, artificial intelligence
What problem does this paper attempt to address?