Med-DANet V2: A Flexible Dynamic Architecture for Efficient Medical Volumetric Segmentation

Haoran Shen1,   Yifu Zhang1,   Wenxuan Wang1,   Chen Chen2,   Jing Liu3,   Shanshan Song1,   Jiangyun Li1,  
1University of Science and Technology Beijing 2University of Central Florida 3Chinese Academy of Sciences
Paper Code arXiv

Abstract

Recent works have shown that the computational efficiency of 3D medical image (e.g. CT and MRI) segmentation can be impressively improved by dynamic inference based on slice-wise complexity. As a pioneering work, a dynamic architecture network for medical volumetric segmentation (i.e. Med-DANet) has achieved a favorable accuracy and efficiency trade-off by dynamically selecting a suitable 2D candidate model from the pre-defined model bank for different slices. However, the issues of incomplete data analysis, high training costs, and the two-stage pipeline in Med-DANet require further improvement. To this end, this paper further explores a unified formulation of the dynamic inference framework from the perspective of both the data itself and the model structure. For each slice of the input volume, our proposed method dynamically selects an important foreground region for segmentation based on the policy generated by our Decision Network and Crop Position Network. Besides, we propose to insert a stage-wise quantization selector to the employed segmentation model (e.g. U-Net) for dynamic architecture adapting. Extensive experiments on BraTS 2019 and 2020 show that our method achieves comparable or better performance than previous state-of-the-art methods with much less model complexity. Compared with previous methods Med-DANet and TransBTS with dynamic and static architecture respectively, our framework improves the model efficiency by up to nearly 4.1 and 17.3 times with comparable segmentation results on BraTS 2019.

Method

MY ALT TEXT

The illustration of the our Med-DANet V2. The Policy Network takes a 2D image slice as input and generates a choice depending on the segmentation difficulty of the current slice. Based on the optimal choice made by the Policy Network, our method can adaptively decide whether to skip the current slice (i.e. directly generating the result with all zero -- "background" class) or send the input with a suitable resolution to the Dynamic Quantization Model (with the cropping region determined by the Crop Position Network). A stage-wise Decision Module is inserted into the Dynamic Quantization Model for suitable model capacity selection. By unifying data-architecture dynamic inference, our method can achieve accurate and efficient segmentation.

Quantitative Results

MY ALT TEXT

To comprehensively evaluate our approach, we conduct experiments on the BraTS 2019 and BraTS 2020 validation set. As shown in the Tables, our method achieves comparable or higher performance than previous SOTA methods with significantly less computational complexity.

Qualitative Results

MY ALT TEXT

A visual comparison of various methods is shown in figure, including 3D U-Net, V-Net, Attention U-Net and the proposed Med-DANet V2. The visual results clearly demonstrate that our framework enhances the delineation of brain tumors, yielding improved segmentation masks by focusing on the tumor regions that are more worth segmenting with the obtained optimal decision in terms of input resolution and quantization.

BibTeX

@InProceedings{Shen_2024_WACV,
    author    = {Shen, Haoran and Zhang, Yifu and Wang, Wenxuan and Chen, Chen and Liu, Jing and Song, Shanshan and Li, Jiangyun},
    title     = {Med-DANet V2: A Flexible Dynamic Architecture for Efficient Medical Volumetric Segmentation},
    booktitle = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)},
    month     = {January},
    year      = {2024},
    pages     = {7871-7881}
}