Shape-intensity knowledge distillation for robust medical image segmentation

TL;DR

Despite its simplicity, experiments on five medical image segmentation tasks of different modalities demonstrate that the proposed Shape-Intensity Knowledge Distillation (SIKD) consistently improves several baseline models under intra-dataset evaluation, and significantly improves the cross-dataset generalization ability.

Abstract

Many medical image segmentation methods have achieved impressive results. Yet, most existing methods do not take into account the shape-intensity prior information. This may lead to implausible segmentation results, in particular for images of unseen datasets. In this paper, we propose a novel approach to incorporate joint shape-intensity prior information into the segmentation network. Specifically, we first train a segmentation network (regarded as the teacher network) on class-wise averaged training images to extract valuable shape-intensity information, which is then transferred to a student segmentation network with the same network architecture as the teacher via knowledge distillation. In this way, the student network regarded as the final segmentation model can effectively integrate the shape-intensity prior information, yielding more accurate segmentation results. Despite its simplicity, experiments on five medical image segmentation tasks of different modalities demonstrate that the proposed Shape-Intensity Knowledge Distillation (SIKD) consistently improves several baseline models (including recent MaxStyle and SAMed) under intra-dataset evaluation, and significantly improves the cross-dataset generalization ability. The code is available at https://github.com/whdong-whu/SIKD.

Figures (8)

• Figure 1: Comparison of averaged performance between the proposed SIKD and corresponding baseline models (including U-Net, SAUNet, PraNet, SANet, TransUnet, MaxStyle, SAMed, LM-Net and 2D D-LKA) under intra-dataset and cross-dataset evaluation on five medical image segmentation tasks of different modalities. The result is the average value across all baseline methods and the corresponding SIKD.
• Figure 2: The pipeline of the proposed method. The teacher and student model have the same network architecture, trained on class-wise averaged training images and original images with segmentation loss, respectively. For the student model, we also apply the distillation loss on the penultimate layer between the teacher and student model to transfer the shape-intensity knowledge.
• Figure 3: Distribution of pixel intensity within each class (e.g., RV cavity, Myocardium, and LV cavity) of the original images (a) and class-wise averaged images (b) on the ACDC training dataset bernard2018deep.
• Figure 4: Some results of the proposed SIKD built upon the baseline U-Net on the cardiac segmentation (Top two rows: intra- and cross-dataset) and some qualitative illustration of SIKD built on the baseline TransUNet for multi-organ segmentation (Bottom two rows: intra- and cross-dataset).
• Figure 5: Some results on the intra-dataset (left two columns) and cross-dataset (right two columns) of polyp segmentation, ONH segmentation, and breast tumor segmentation (from top to bottom). Green outline: segmentation by the baseline U-Net model; Blue outline: segmentation by SIKD built upon U-Net; Light blue area: ground-truth segmentation.