2229
Uncertainty-guided task-specific multi-parametric MR image fusion for brain tissue segmentation and quantification1Paul C. Lauterbur Research Center for Biomedical Imaging, Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences, Shenzhen, China, 2University of Chinese Academy of Sciences, Beijing, China, 3Peng Cheng Laboratory, Shenzhen, China
Synopsis
Keywords: Analysis/Processing, Brain
Motivation: Existing techniques for multi-parametric MR imaging-based brain tissue segmentation typically employ a generic feature combination strategy without incorporating task-specific guidance, making it challenging to ensure effective fusion.
Goal(s): In this work, we aim to develop a task-specific multi-parametric MR image fusion framework to enhance the brain tissue segmentation and quantification accuracy.
Approach: During preliminary experiments, we have identified a close correlation between prediction uncertainties and prediction errors. Therefore, we propose an uncertainty-guided task-specific multi-parametric MR image fusion framework to enhance fusion efficiency and decrease prediction uncertainty.
Results: Experiments on the iSeg-2019 dataset demonstrate that the proposed method achieves better results than existing techniques.
Impact: The outcome of this work has the potential to be utilized in clinical practice to help physicians better monitor brain development and diagnose brain diseases. Meanwhile, the framework can be extended to diverse fields where multi-modal image fusion is required.
Introduction
Brain tissue segmentation plays a crucial role in quantifying brain structure volumes, investigating infant brain development patterns, and assessing brain disease progression 1. Multi-parametric MR imaging serves as a valuable tool in achieving these objectives. However, due to issues such as low image contrast and complex anatomical structures, manually segmenting brain tissues, including white matter (WM), gray matter (GM), and cerebrospinal fluid (CSF), in multi-parametric MR images is difficult and time-consuming 2. As a result, significant attention has been devoted to the development of automatic algorithms.Recently, many deep learning-based methods have been developed to segment brain tissues in multi-parametric MR images 2–4. For example, Joze et al. proposed the HyperDense-Net, which utilizes dense connections to aggregate the multi-parametric information 4. In general, there exist three major types of deep learning-based multi-parametric image fusion methods: early fusion 5,6, late fusion 7,8, and multi-layer fusion 4. Although it is commonly believed that multi-layer fusion yields the best performance, existing methods tend to fuse the extracted image features directly through concatenation or summation without incorporating task-specific guidance, making it intractable to guarantee that the fusion is optimal for the specific task in hand.
In this study, we propose an uncertainty-guided task-specific multi-parametric MR image fusion framework for brain tissue segmentation. Our framework is built upon the observation that the prediction errors of one network are closely correlated with its prediction uncertainties (Figure 1). Accordingly, we hypothesize that we can enhance the network performance by controlling the uncertainties in network predictions. To achieve this, our framework incorporates two key contributions. Firstly, we propose a task-specific multi-parametric MR image fusion approach guided by uncertainty (M-SUM in Figure 1). Secondly, we design an uncertainty-enhanced loss function to directly minimize the uncertainty in segmentation results ($$$L_{AUG}$$$ in Figure 1). Experimental results demonstrate that our proposed framework can effectively improve brain tissue segmentation performance.
Methods
The architecture of our proposed framework is depicted in Figure 2. Our framework consists of multiple streams to extract features from MRI data acquired with different sequences. Briefly, four augmented inputs are generated from the same image data. Three are utilized to generate three sets of fused outputs ($$$\widehat{p_f^a}$$$ ($$$a\in\{2,3,4\}$$$) in Figure 2), while the remaining one is utilized to generate the modality-specific segmentation probability outputs ($$$\widehat{p_{m1}^1}$$$ and $$$\widehat{p_{m2}^1}$$$) and the final fused output ($$$\widehat{p_f^1}$$$).In our proposed module, M-SUM, two modality-specific uncertainty maps are calculated from $$$\widehat{p_{m1}^1}$$$ and $$$\widehat{p_{m2}^1}$$$:
$$U_{m1}=exp(1-\frac{\max_c\widehat{p_{m1}^1}}{\min_c\widehat{p_{m1}^1}})$$
$$U_{m2}=exp(1-\frac{\max_c\widehat{p_{m2}^1}}{\min_c\widehat{p_{m2}^1}})$$
Where $$$c$$$ represents the segmentation category. Then, these two uncertainty maps are utilized to optimize the features extracted from different imaging sequences. Specifically, with the modality-specific uncertainty maps, three sets of modality-specific features are obtained and aggregated according to the procedure depicted in Figure 3.
The whole framework is trained using four losses. Among them, three ($$$L_{seg,m1}$$$, $$$L_{seg,m2}$$$, and $$$L_{seg,f}$$$) adopt the classic cross-entropy loss, while the last one is the specially designed loss, $$$L_{AUG}$$$. Here, an epistemic uncertainty $$$U_{fo}$$$ is firstly estimated 9,10:
$$U_{fo}\approx\frac{1}{3}\sum_{r=2}^{4}(\widehat{p_f^r})^2-(\frac{1}{3}\sum_{r=2}^{4}\widehat{p_f^r})^2$$
We then compress $$$U_{fo}$$$ by keeping the largest uncertainty values among the different categories and normalize it into a reasonable range of [0,1], obtaining $$$U_f$$$.
$$$L_{AUG}$$$ is defined as:
$$L_{AUG}=-\frac{1}{H \times W \times D}\sum_{h=1}^{H}\sum_{w=1}^{W}\sum_{d=1}^{D}U_f(h,w,d)$$
Finally, the total training loss is:
$$L=\alpha \cdot (L_{seg,m1}+L_{seg,m2})+\beta \cdot L_{seg,f}+\gamma \cdot L_{AUG}$$
Where $$$\alpha$$$, $$$\beta$$$, and $$$\gamma$$$ are three weighting parameters determined empirically.
Results and Discussion
All our experiments were conducted using the open-source dataset, iSeg-2019, which provides 10 3D paired T1-weighted and T2-weighted MR images (train:val:test=6:1:3). To quantitatively evaluate the segmentation performance, three metrics are calculated and reported. Table 1 lists the results of different methods. Three observations can be made: 1) Our baseline model (without M-SUM and $$$L_{AUG}$$$) can already achieve better segmentation results than the existing method, HyperDense-Net. 2) Both the two components in our framework, M-SUM and $$$L_{AUG}$$$, contribute to enhancing the segmentation performance. 3) Our final model achieves the best results.Figure 4 visualizes two example segmentation results of HyperDense-Net and our final model. It can be clearly observed that our method achieves much better performance, especially for the segmentation of GM and WM, validating the effectiveness of the proposed uncertainty-guided task-specific multi-parametric MR image fusion approach.
Conclusion
In this study, an uncertainty-guided task-specific multi-parametric MR image fusion network is developed for brain tissue segmentation. Two key contributions are proposed, including a modality-specific uncertainty-aware feature fusion module and an augmentation ensembling uncertainty-enhanced loss function. Experimental results demonstrate that the proposed method effectively enhance the segmentation performance of three brain tissues, which can be very helpful in clinical practice to help physicians better monitor brain development and diagnose brain diseases.Acknowledgements
This research was partly supported by the National Natural Science Foundation of China (62222118, U22A2040), Guangdong Provincial Key Laboratory of Artificial Intelligence in Medical Image Analysis and Application (2022B1212010011), Shenzhen Science and Technology Program (RCYX20210706092104034, JCYJ20220531100213029), and Key Laboratory for Magnetic Resonance and Multimodality Imaging of Guangdong Province (2023B1212060052).References
1. Sun Y, Gao K, Wu Z, et al. Multi-site infant brain segmentation algorithms: The iSeg-2019 challenge. IEEE Trans. Med. Imaging. 2021; 40(5): 1363–1376.
2. Zhuang Y, Liu H, Song E, et al. APRNet: A 3D anisotropic pyramidal reversible network with multi-modal cross-dimension attention for brain tissue segmentation in MR images. IEEE J. Biomed. Heal. Informatics. 2022; 26(2): 749–761.
3. Li J, Yu ZL, Gu Z, et al. MMAN: Multi-modality aggregation network for brain segmentation from MR images. Neurocomputing. 2019; 358: 10–19.
4. Dolz J, Gopinath K, Yuan J, et al. HyperDense-Net: A hyper-densely connected CNN for multi-modal image segmentation. IEEE Trans. Med. Imaging. 2019; 38(5): 1116–1126.
5. Zhou C, Ding C, Lu Z, et al. One-pass multi-task convolutional neural networks for efficient brain tumor segmentation. In MICCAI 11072, 637–645 (Springer International Publishing, 2018).
6. Zhang W, Li R, Deng H, et al. Deep convolutional neural networks for multi-modality isointense infant brain image segmentation. Neuroimage. 2015; 108: 214–224.
7. Pinto A, Pereira S, Meier R, et al. Enhancing clinical MRI perfusion maps with data-driven maps of complementary nature for lesion outcome prediction. In MICCAI 11072, 107–115 (Springer International Publishing, 2018).
8. Nie D, Wang L, Gao Y, et al. Fully convolutional networks for multi-modality isointense infant brain image segmentation. In IEEE ISBI, 1342–1345 (IEEE, 2016).
9. Kendall A, Gal Y. What uncertainties do we need in Bayesian deep learning for computer vision? In NIPS (2017).
10. Cao X, Chen H, Li Y, et al. Uncertainty aware temporal-ensembling model for semi-supervised ABUS mass segmentation. IEEE Trans. Med. Imaging. 2021; 40(1): 431–443.
Figures
