A Semi-Supervised Abdominal Multi-Organ Pan-Cancer Segmentation Framework with Knowledge Distillation and Multi-Label Fusion

Zengmin Zhang, Xiaomeng Duan, Yanjun Peng, Zhengyu Li

Keywords: Knowledge Distillation, Dual Attention, Multiple labels fusion, Semi-supervised learning

TL;DR: In this paper, we propose a semi-supervised knowledge distillation framework that achieves efficient segmentation while maintaining accuracy.

Abstract: The segmentation of abdominal organs and tumors plays a crucial role in computer-aided diagnosis of medical images. To achieve high-precision segmentation while maintaining efficiency, especially in semi-supervised learning, we propose a novel semi-supervised knowledge distillation framework. The framework consists of the teacher model and the student model. In the first step, we design an attention nnU-Net with a dual convolutional attention decoder as the teacher model to generate high-quality tumor pseudo-labels for unlabeled tumor data. The dual attention decoder enhances attention to the regions of interest and highlights the most relevant channels, improving the model's ability to optimize features. Additionally, we design an effective 2D sliding window inference strategy to accelerate the inference speed of the teacher model. We utilize partial labels, organ pseudo-labels provided by the FLARE2022 winner, and tumor pseudo-labels for multi-label fusion, ensuring the fusion results closely resemble the ground truth. In the second step, we employ a lightweight nnU-Net as the student model to achieve efficient segmentation. Our method achieved an average DSC score of 88.53\% and 30.47\% for the organs and lesions on the validation set and the average running time and area under GPU memory-time cure are 15.85s and 15601MB, respectively. Our code is available at https://github.com/zzm3zz/FLARE2023.

Submission Number: 9