Efficient Multi-Organ Segmentation From 3D Abdominal CT Images With Lightweight Network and Knowledge Distillation
Qianfei Zhao, Lanfeng Zhong, Jianghong Xiao, Jingbo Zhang, Yinan Chen, Wenjun Liao, Shaoting Zhang, Guotai Wang
Abstract
Accurate segmentation of multiple abdominal organs from Computed Tomography (CT) images plays an important role in computer-aided diagnosis, treatment planning and follow-up. Currently, 3D Convolution Neural Networks (CNN) have achieved promising performance for automatic medical image segmentation tasks. However, most existing 3D CNNs have a large set of parameters and huge floating point operations (FLOPs), and 3D CT volumes have a large size, leading to high computational cost, which limits their clinical application. To tackle this issue, we propose a novel framework based on lightweight network and Knowledge Distillation (KD) for delineating multiple organs from 3D CT volumes. We first propose a novel lightweight medical image segmentation network named LCOV-Net for reducing the model size and then introduce two knowledge distillation modules (i.e., Class-Affinity KD and Multi-Scale KD) to effectively distill the knowledge from a heavy-weight teacher model to improve LCOV-Net's segmentation accuracy. Experiments on two public abdominal CT datasets for multiple organ segmentation showed that: 1) Our LCOV-Net outperformed existing lightweight 3D segmentation models in both computational cost and accuracy; 2) The proposed KD strategy effectively improved the performance of the lightweight network, and it outperformed existing KD methods; 3) Combining the proposed LCOV-Net and KD strategy, our framework achieved better performance than the state-of-the-art 3D nnU-Net with only one-fifth parameters. The code is available at https://github.com/HiLab-git/LCOVNet-and-KD.